[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/design/feature/offloader/distributed_layerwise_offload.md

@@ -19,15 +19,17 @@ DLO weight collective.
 
 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.
+otherwise it uses the ordinary loader. In no-AllGather mode, an opt-in Phase B
+runtime cache can normalize those final ordinary-loader DiT tensors into an
+immutable node-local mmap entry. Consequently, replicas can share either
+proven-compatible checkpoint pages or equivalent final runtime layouts. An
+opt-in CUDA registration budget can make the complete final mapping directly
+H2D-copyable, avoiding the recurrent host packing otherwise required by mmap.
 
 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.
+tensors. Direct checkpoint mmap still does not cover TP, but the Phase B cache
+can share each final TP coordinate among equivalent processes.
 
 The compatibility matrix below describes the current implementation. The
 unit-level guards are covered, but not every parallelism combination has a
@@ -81,6 +83,90 @@ 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.
 
+### Post-load runtime cache
+
+The Phase B cache preserves the same ownership boundary. DLO neither derives
+cache keys nor interprets model loaders. When enabled with no-AllGather, every
+rank first completes the ordinary loader, custom callbacks, post-load
+processing, strict validation, and calibration. The loader then:
+
+1. discovers final CPU DiT parameters and persistent buffers;
+2. rejects unsupported non-contiguous, aliased, quantized, distributed, or
+   ambiguous layouts before changing the model;
+3. hashes tensor names, metadata, and final bytes;
+4. builds or joins one content-addressed entry under a strict POSIX lock; and
+5. returns a `HostWeightPlan(backing_kind="runtime_cache")` only after file,
+   manifest, and remapped-tensor content validation succeeds.
+
+The identity excludes DP size/rank and SP rank. It includes TP world size/rank
+and conservative SP implementation/world-size guards. DP replicas and
+equivalent SP ranks can therefore share; different TP coordinates cannot.
+Final tensor content is authoritative, so an unmodeled loader/config change
+that changes bytes creates a different entry instead of risking silent
+sharing.
+
+Publication writes sharded safetensors and a JSON manifest into a same-
+filesystem temporary directory, flushes them, and atomically renames the
+directory. The manifest includes per-file SHA-256 digests and the final tensor
+content digest. Kernel `flock` ownership distinguishes a live writer from a
+stale lock file: process death releases the lock automatically. Stale temporary
+entries are removed only by the next lock owner. A bounded wait or any cache
+error retains the already-valid ordinary tensors.
+
+DLO maps this plan without moving the DiT to `meta` and without rerunning any
+post-load hook. It changes tensor storage while preserving `Parameter` objects,
+then runs the existing `gc.collect()`/`malloc_trim()` cleanup before installing
+streaming hooks. The source mappings remain open through DLO disable. With the
+default zero registration budget they feed the existing bounded two-slot host
+staging path. Every streamed block is first packed from pageable mmap storage
+into one of those pinned slots and is then copied to the device. This removes
+model-sized private pinned storage, but adds a recurrent host-to-host copy and
+can amplify TP synchronization waits when several ranks compete for host
+bandwidth.
+
+This version is deliberately opt-in and steady-state-oriented. It requires
+roughly one local-disk model copy per runtime identity, has no eviction, and
+still performs ordinary loading plus full content-hash passes in every rank.
+Skip-load-on-hit is a separate follow-up after loader side effects can be
+modeled safely.
+
+### Registered runtime-cache transfer
+
+`dlo_runtime_cache_pin_limit_gib > 0` opts a CUDA worker into complete
+runtime-cache registration. It also changes loader selection: no-AllGather
+uses the ordinary loader and final runtime cache even when a TP1 direct-
+checkpoint plan is available. A checkpoint binding can carry a deferred
+adapter such as grouped-QKV reordering. Registering that raw source would not
+remove the adapter's recurrent CPU allocation/copy, whereas the runtime cache
+contains the transform-complete final tensors.
+
+After mapping and validating every final tensor, DLO groups source storages by
+backing file, page-aligns their address ranges, and coalesces overlapping or
+adjacent ranges within that mapping. The complete byte count is checked
+against the per-worker budget before any CUDA call. It then registers every
+range with the read-only CUDA host-registration flag and verifies that PyTorch
+recognizes every mapped tensor as pinned. Failure rolls back any ranges already
+registered and preserves the existing two-slot staging path. Partial direct
+transfer is deliberately unsupported. A rollback failure aborts worker
+initialization so a still-registered range is retained until process/context
+teardown rather than being unsafely unmapped.
+
+On success, the hooks copy each immutable tensor view directly into its offset
+in the existing flattened rotating device buffer on the copy stream. Parameter
+repointing, two-block HBM residency, and H2D/compute overlap are unchanged. The
+implementation currently issues more H2D copy operations than block packing,
+but avoids the much larger recurrent CPU copy. A packed cache schema should be
+considered only if launch-count measurements justify its added format and
+publication complexity.
+
+Registration is per process and CUDA context; physical file-cache pages remain
+shared and are not copied into one private host allocation per worker. The
+operator must nevertheless budget page-locked memory and satisfy OS/CUDA
+registration limits for every worker. During shutdown, DLO first synchronizes
+outstanding device work, unregisters all ranges, and then closes the safetensors
+mappings. Worker teardown invokes this cleanup before destroying the
+distributed environment and CUDA context.
+
 ### AllGather path
 
 With the default `dlo_use_allgather=True`, each rank stores approximately
@@ -114,11 +200,17 @@ 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.
 
+Fast NVLink/NVSwitch changes collective cost, not this execution contract. For
+independently scheduled replicas—even on one NVSwitch baseboard—no-AllGather
+plus shared mmap storage avoids coupling their request schedules and failure
+domains.
+
 ### Rank-local path without DLO AllGather
 
 With `--dlo-no-use-allgather`, DLO forces its internal offload shard size to
 one and streams complete blocks using H2D copies only. The host backing may be
-either a loader-approved checkpoint mapping or ordinary runtime tensors.
+a loader-approved checkpoint mapping, a normalized runtime-cache mapping, or
+ordinary runtime tensors.
 
 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
@@ -127,10 +219,14 @@ 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.
 
+For a runtime-cache mapping, a positive registration budget replaces those
+host slots with direct copies from the shared final tensor views. A zero budget
+or failed registration retains the bounded staging behavior.
+
 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.
+parameters. Supported CPU layouts may then enter the opt-in runtime cache.
+Unsupported layouts keep one private runtime copy per process.
 
 This mode means:
 
@@ -141,17 +237,17 @@ This mode means:
 - TP/HSDP/SP collectives, if configured, are not disabled by this flag; only
   DLO's additional weight AllGather is disabled.
 - Pure DP deployments share one checkpoint-backed copy per node when direct
-  mmap is selected; the ordinary-loader fallback keeps one private runtime
-  copy per rank.
+  mmap is selected. The runtime cache can share equivalent ordinary-loader
+  outputs; its fallback keeps one private runtime copy per rank.
 - The scheduler does not require a synchronized DP request wave for DLO.
 
 ## 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. 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** | 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. |
+| **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; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. CUDA workers may register the complete final mapping for direct H2D. |
+| **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 without a DLO weight collective. Runtime-cache v1 excludes SP rank and includes conservative SP implementation/world-size guards. Registered transfer does not remove ordinary SP activation/attention collectives. |
+| **TP > 1** | Outside the Phase A direct-mmap scope. The loader preserves TP-local layouts and DLO may apply DP/SP host sharding to ordinary runtime tensors. | The ordinary TP-aware loader produces rank-local tensors. Runtime-cache v1 creates one entry per TP coordinate, shareable by equivalent DP/SP processes. Registering those final mappings removes staging-driven TP skew but does not remove ordinary TP collectives. |
 | **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
@@ -182,12 +278,10 @@ requirements.
 Direct checkpoint mmap can back either transfer path. It is currently limited
 to proven TP1, non-HSDP, non-online-quantized layouts. Other layouts use the
 ordinary loader. Online quantization remains incompatible with DLO AllGather;
-use `--dlo-no-use-allgather` or disable online quantization.
-
-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).
+use `--dlo-no-use-allgather` or disable online quantization. Runtime-cache v1
+is no-AllGather only and rejects all quantization configurations, HSDP/DTensor,
+expert parallelism, CFG parallelism, PP, non-contiguous tensors, and
+tied/shared storage.
 
 ## Validation coverage
 
@@ -200,6 +294,13 @@ Current source-level validation includes:
 - 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;
+- runtime-cache content identity, sharded atomic publication, concurrent
+  writer election, corruption rebuild, lock timeout, and fail-closed layout
+  rejection;
+- loader ordering through final post-load mutation and DLO remapping without
+  rerunning post-load hooks;
+- all-or-nothing page-aligned registration, budget rejection, partial-failure
+  rollback, direct-copy staging bypass, and explicit shutdown cleanup;
 - resident-layer requests requiring no-AllGather;
 - DP request-wave validation for denoising-step compatibility;
 - sharding, double-buffer, AllGather-size, and heterogeneous-block regression
@@ -259,11 +360,99 @@ from 167.53 to 70.24 GiB for worker 0 and from 115.40 to 17.44 GiB for worker
 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.
+### L20X runtime-cache smoke matrix
+
+Phase B was also exercised on four L20X GPUs in one NVLink domain with
+MiniMax-H3 FL2VA, CUDNN attention, eager execution, 256x256 four-second
+outputs, two denoising steps, `dlo_resident_layers=0`, text-encoder TP1, and
+VAE TP1. This is a functional, memory, and profiler smoke test rather than a
+production-quality performance benchmark. Host values below are node-wide
+maxima during request waves from descendant-process `smaps_rollup`; they do
+not represent the private startup peak. HBM is the engine-reported request
+peak because sampled `nvidia-smi` data can miss short-lived allocations.
+
+One internal DP2xTP2 engine compared the three host-storage/transfer choices.
+Each measured wave contained two requests. NVLink TX is the physical per-link
+counter delta summed across the four GPUs; RX had the same value and is not
+added again here.
+
+| Mode | Two-request wave | Throughput | Host PSS | `Private_Dirty` | Engine peak HBM | NVLink TX / wave |
+|---|---:|---:|---:|---:|---:|---:|
+| AllGather, ordinary loader | 3.41 s | 0.587 req/s | 210.70 GiB | 209.72 GiB | 12,498 MB | 73.37 GiB |
+| no-AllGather, private loader weights | 6.96 s | 0.287 req/s | 312.56 GiB | 311.56 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (cold publish) | 18.94 s | 0.106 req/s | 178.71 GiB | 116.00 GiB | 11,882 MB | 19.78 GiB |
+| no-AllGather, runtime cache (warm reuse) | 24.23 s | 0.083 req/s | 178.29 GiB | 115.58 GiB | 11,882 MB | 19.78 GiB |
+
+The no-AllGather cache preserved the no-AllGather NVLink traffic and HBM
+footprint while reducing node PSS by about 134 GiB relative to private loader
+weights. Cold and warm labels describe cache publication or reuse during
+initialization; both request phases ran after initialization. A warm hit still
+performed ordinary loading and full content hashing, so it was not a startup
+fast path. All measured no-AllGather private/cache outputs were byte-identical
+for matching seeds.
+
+The deployment shape targeted by the cache was tested separately as two
+independent DP1xTP2 engines on the same node. Both engines mapped the same two
+TP-coordinate entries, including the same file device/inode tuples; the cache
+occupied 66,281,995,903 bytes. A start barrier kept all four measured requests
+ahead of profiling.
+
+| Storage | Mean request latency, engines A / B | Combined throughput | Host PSS | `Private_Dirty` | Engine peak HBM |
+|---|---:|---:|---:|---:|---:|
+| Private loader weights | 3.50 / 3.60 s | 0.463 req/s | 366.20 GiB | 365.22 GiB | 11,882 MB |
+| Shared runtime cache | 8.96 / 12.30 s | 0.142 req/s | 232.23 GiB | 169.52 GiB | 11,882 MB |
+
+The shared cache reduced request-time PSS by 133.97 GiB (36.6%) and
+`Private_Dirty` by 195.70 GiB (53.6%), but reduced combined throughput by
+69.3% in this host-contention smoke. Across one profiled request per engine,
+aggregate GPU compute was unchanged (1.359 versus 1.357 seconds) and H2D
+payload was exactly 264.24 GiB in both modes. Aggregate CPU `aten::copy_` time
+rose from 3.75 to 32.49 seconds, while communication-kernel time rose from
+3.56 to 5.80 seconds as TP ranks waited on skewed staging. These durations are
+summed across ranks and can overlap; they identify the mechanism, not wall
+time. Matching private/cache outputs were byte-identical.
+
+The result confirms that zero-budget runtime-cache staging is a host-capacity
+option, not a latency-neutral replacement for private pinned weights. Broader
+SP, model, platform, and production-quality validation remains tracked in
+[issue #6231](https://github.com/vllm-project/vllm-omni/issues/6231). The
+historical B300 rows above predate the runtime cache and must not be treated as
+Phase B memory validation.
+
+### L20X registered-runtime-cache smoke
+
+The read-only registered path was then tested with the same MiniMax-H3
+workload. These are functional and mechanism-validation measurements, not a
+production benchmark. The TP1 pair used one request per wave. Both modes moved
+exactly 121.86 GiB H2D in the profiled request and produced byte-identical video
+and audio for both measured seeds.
+
+| TP1 no-AllGather mode | Mean wave | Throughput | Request-time host PSS | `Private_Dirty` | Engine peak HBM | CPU `aten::copy_` | GPU compute |
+|---|---:|---:|---:|---:|---:|---:|---:|
+| Checkpoint mmap + bounded staging | 33.87 s | 0.0295 req/s | 133.65 GiB | 71.16 GiB | 13,226 MB | 21.91 s | 419.28 ms |
+| Final runtime cache + registered direct H2D | 4.88 s | 0.2055 req/s | 129.41 GiB | 66.93 GiB | 13,226 MB | 1.64 s | 419.05 ms |
+
+Registration covered 61.73 GiB in 13 page-aligned ranges. Mean wave latency
+fell by 85.6% and throughput increased 6.94x. GPU compute and H2D payload were
+unchanged, identifying recurrent host transform/packing as the removed work.
+The direct path issued 2,425 H2D events versus 1,945 for packed staging, but
+that launch increase was much smaller than the eliminated CPU cost.
+
+Two independent TP2 engines also registered the same two TP-coordinate cache
+entries across four workers. Compared with the zero-budget cache, mean wave
+latency improved from 9.21/12.55 seconds to 3.80/3.86 seconds for engines A/B;
+the private-loader reference was 3.61/3.72 seconds. Registered request-time
+host PSS was 224.26 GiB versus 232.23 GiB staged and 366.20 GiB private.
+Aggregate CPU `aten::copy_` time was 3.83 seconds registered, 32.49 seconds
+staged, and 3.75 seconds private; aggregate GPU compute and the 264.24 GiB H2D
+payload were effectively unchanged. Ordinary TP communication remained, but
+the extra wait caused by staging-driven rank skew largely disappeared.
+
+All workers explicitly unregistered before releasing their mappings, the
+system's acquired-minus-released pin counter returned to its pre-run value,
+and every GPU allocation was released. Cold runtime-cache publication remained
+expensive because it still performs ordinary loading and multiple full-content
+passes; that startup problem is outside the registered transfer scope.
 
 ## Recommendations
 
@@ -272,8 +461,12 @@ change the Phase A direct-mmap TP1 support boundary.
 - Use **SP + DLO AllGather** for long-sequence workloads when DP concurrency is
   not the goal.
 - Use **no-AllGather** when independent replica execution is required. TP1
-  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.
+  direct-mmap deployments can share checkpoint pages per node. Enable the
+  runtime cache only when independent replicas have repeated final layouts and
+  benchmark the target CPU/memory topology, especially for matching
+  coordinates across multiple TP engines. On CUDA, provide an explicit
+  registration budget large enough for the complete final layout when
+  recurrent staging cost is unacceptable; leave it at zero when page-locking
+  that layout is not operationally acceptable.
 - 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

@@ -5,7 +5,11 @@ 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 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.
+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. 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.
@@ -21,8 +25,10 @@ AllGather:       [AG N+1]           [AG N+2]
 Slots:           slot 0: Layer N    slot 1: Layer N+1
 ```
 
-AllGather communicates only request-independent weight shards, so data-
-parallel ranks may process different requests concurrently.
+AllGather communicates request-independent weight shards, but every member of
+the transfer group must request the same next block in the same collective
+order. A synchronized DP wave may contain different prompts only when those
+requests follow the same denoising and block-execution path.
 
 ## Usage
 
@@ -38,6 +44,15 @@ vllm serve /path/to/model --omni \
   --data-parallel-size 4 \
   --dlo-no-use-allgather
 
+# Independently scheduled replicas sharing final runtime weights on one node
+vllm serve /path/to/model --omni \
+  --enable-distributed-layerwise-offload \
+  --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-pin-limit-gib 80
+
 # Sequence parallel deployment
 vllm serve /path/to/model --omni \
   --enable-distributed-layerwise-offload \
@@ -50,7 +65,10 @@ from vllm_omni import Omni
 omni = Omni(
     model="/path/to/model",
     enable_distributed_layerwise_offload=True,
-    dlo_use_allgather=True,
+    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,
 )
 ```
 
@@ -62,6 +80,10 @@ omni = Omni(
 | `--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` | Stream complete rank-local blocks without a DLO weight collective | `false` |
+| `--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
@@ -78,6 +100,50 @@ 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, 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
+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.
+
+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
 use their ordinary component loader. A checkpoint source that mixes DiT and
@@ -104,6 +170,34 @@ 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.
 
+## Runtime-cache lifecycle
+
+Runtime-cache publication uses a per-layout POSIX file lock, a temporary
+directory on the same filesystem, file and final-tensor SHA-256 validation,
+`fsync`, and an atomic directory rename. A crashed writer releases its kernel
+lock; the next writer removes only that layout's stale temporary directories
+while holding the lock. Lock timeout, unsupported layout, insufficient disk,
+or validation failure keeps the already-valid ordinary tensors.
+
+The cache root is deliberately independent of `VLLM_CACHE_ROOT`, because
+vLLM-Omni may isolate that variable between stage replicas. All consumers that
+should share must receive the same `--dlo-runtime-cache-dir` on a local,
+disk-backed filesystem. Do not place the cache on tmpfs when the goal is lower
+host-memory PSS, and do not use a cross-node directory in this version.
+
+There is no automatic eviction. Budget roughly one final DiT copy per distinct
+TP coordinate, dtype/layout, and model version, plus temporary space while an
+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:
@@ -132,22 +226,37 @@ must enter each collective.
 
 ## Limitations
 
-- Direct checkpoint mmap currently requires TP1. TP greater than one is
-  outside the Phase A shared-mmap support scope and falls back before model
-  mutation to the ordinary TP-aware loader. DLO can stream that runtime layout,
-  but it provides no shared-mmap host-memory benefit or guarantee.
+- 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. 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.
 - Online quantization uses the ordinary loader with no-AllGather. It remains
-  incompatible with DLO AllGather.
+  incompatible with DLO AllGather and runtime-cache v1. Runtime-cache v1 also
+  rejects other quantized, non-contiguous, aliased/tied, device-only, HSDP,
+  expert-parallel, CFG-parallel, and PP layouts.
 - 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.
+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. 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).
+
+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%. 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.

tests/diffusion/model_loader/test_diffusers_loader.py

@@ -12,6 +12,7 @@ import torch
 import torch.nn as nn
 from huggingface_hub import snapshot_download
 from vllm.config.load import LoadConfig
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 
 from vllm_omni.diffusion.config import get_current_diffusion_config, get_current_diffusion_config_or_none
 from vllm_omni.diffusion.data import OmniDiffusionConfig
@@ -335,7 +336,7 @@ def test_dlo_plan_loads_component_sources_outside_planned_dit(monkeypatch):
         backing_kind="checkpoint_mmap",
         bindings={
             "transformer.weight": TensorBinding(
-                checkpoint_key="weight",
+                storage_key="weight",
                 file_path="unused",
             )
         },
@@ -394,6 +395,192 @@ def test_dlo_plan_fallback_runs_ordinary_loader(monkeypatch):
     assert loader.take_host_weight_plan() is None
 
 
+def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+            ulysses_degree=1,
+            ring_degree=1,
+            allgather_degree=1,
+            ulysses_mode="strict",
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_lock_timeout=1.0,
+        model="unused",
+        revision=None,
+        model_config={},
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = UnquantizedLinearMethod()
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    def load_weights(_model):
+        events.append("load")
+        with torch.no_grad():
+            model.transformer.weight.fill_(1)
+
+    def process_weights(_model, _device):
+        events.append("process")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    def calibrate(_model):
+        events.append("calibrate")
+        with torch.no_grad():
+            model.transformer.weight.add_(1)
+
+    original_eval = model.eval
+
+    def eval_model():
+        events.append("eval")
+        return original_eval()
+
+    def build_cache(_pipeline, **kwargs):
+        events.append("cache")
+        assert torch.equal(model.transformer.weight, torch.full((2, 2), 3.0))
+        assert kwargs["tensor_parallel_rank"] == 0
+        assert "data_parallel_size" not in kwargs["sequence_parallel_guard"]
+        assert kwargs["quantization_config"] is None
+        return HostWeightPlanResult(runtime_plan)
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = load_weights  # type: ignore[method-assign]
+    loader._process_weights_after_loading = process_weights  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = calibrate  # type: ignore[method-assign]
+    model.eval = eval_model  # type: ignore[method-assign]
+    monkeypatch.setattr(
+        loader_mod,
+        "build_checkpoint_mmap_plan",
+        lambda *_args, **_kwargs: HostWeightPlanResult(None, "requires normalization"),
+    )
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "calibrate", "eval", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_registration_prefers_final_runtime_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=1,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        dlo_runtime_cache_pin_limit_gib=1.0,
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    events: list[str] = []
+    runtime_plan = HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings={},
+        runtime_layout_key="layout-key",
+        post_load_complete=True,
+    )
+
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: events.append("load")  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: events.append("process")  # type: ignore[method-assign]
+    loader._apply_skip_softmax_calibration = lambda _model: None  # type: ignore[method-assign]
+
+    def unexpected_checkpoint_plan(*_args, **_kwargs):
+        pytest.fail("registered runtime-cache mode must not select the raw checkpoint mmap plan")
+
+    def build_cache(_pipeline, **_kwargs):
+        events.append("cache")
+        return HostWeightPlanResult(runtime_plan)
+
+    monkeypatch.setattr(loader_mod, "build_checkpoint_mmap_plan", unexpected_checkpoint_plan)
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    assert loader.load_model(load_device="cpu") is model
+    assert events == ["load", "process", "cache"]
+    assert loader.take_host_weight_plan() is runtime_plan
+
+
+def test_dlo_runtime_cache_failure_retains_ordinary_weights(monkeypatch, tmp_path):
+    import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
+    import vllm_omni.diffusion.model_loader.runtime_weight_cache as cache_mod
+
+    od_config = SimpleNamespace(
+        dtype=torch.float32,
+        parallel_config=SimpleNamespace(
+            use_hsdp=False,
+            tensor_parallel_size=1,
+            sequence_parallel_size=1,
+            data_parallel_size=2,
+            cfg_parallel_size=1,
+            pipeline_parallel_size=1,
+        ),
+        quantization_config=None,
+        enable_distributed_layerwise_offload=True,
+        dlo_use_allgather=False,
+        dlo_enable_runtime_cache=True,
+        dlo_runtime_cache_dir=str(tmp_path),
+        model="unused",
+    )
+    loader = DiffusersPipelineLoader(LoadConfig(), od_config)
+    model = nn.Module()
+    model.transformer = nn.Linear(2, 2, bias=False)
+    model.transformer.quant_method = object()
+    loader._init_from_load_format = lambda *_args, **_kwargs: model  # type: ignore[method-assign]
+    loader.load_weights = lambda _model: None  # type: ignore[method-assign]
+    loader._process_weights_after_loading = lambda *_args: None  # 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, "requires normalization"),
+    )
+
+    def build_cache(_pipeline, **kwargs):
+        assert kwargs["quantization_config"] == "module_quant_method"
+        return HostWeightPlanResult(None, "writer timed out", "lock_timeout")
+
+    monkeypatch.setattr(cache_mod, "build_runtime_weight_cache_plan", build_cache)
+
+    original_weight = model.transformer.weight.detach().clone()
+    assert loader.load_model(load_device="cpu") is model
+    assert torch.equal(model.transformer.weight, original_weight)
+    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/model_loader/test_runtime_weight_cache.py

@@ -0,0 +1,222 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+import fcntl
+import multiprocessing as mp
+import shutil
+from pathlib import Path
+
+import pytest
+import torch
+from torch import nn
+
+from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+    build_runtime_weight_cache_plan,
+    default_runtime_cache_root,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _LoaderMarker:
+    pass
+
+
+class _Transformer(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.proj = nn.Linear(3, 2)
+        self.register_buffer("persistent", torch.arange(2, dtype=torch.float32))
+        self.register_buffer("temporary", torch.ones(1), persistent=False)
+        with torch.no_grad():
+            self.proj.weight.copy_(torch.arange(6, dtype=torch.float32).reshape(2, 3))
+            self.proj.bias.copy_(torch.tensor([7.0, 8.0]))
+
+
+class _Pipeline(nn.Module):
+    def __init__(self):
+        super().__init__()
+        self.transformer = _Transformer()
+
+
+def test_default_root_is_not_split_by_replica_vllm_cache(monkeypatch):
+    monkeypatch.setenv("VLLM_CACHE_ROOT", "/tmp/replica-specific-vllm-cache")
+
+    assert default_runtime_cache_root().endswith("/.cache/vllm-omni/dlo-runtime-weights")
+
+
+def _build(pipeline: nn.Module, cache_root: Path, **overrides):
+    kwargs = {
+        "dit_modules": (("transformer", pipeline.transformer),),
+        "loader_type": _LoaderMarker,
+        "cache_root": cache_root,
+        "lock_timeout_seconds": 2.0,
+        "max_shard_bytes": 32,
+        "model_identity": "test/model",
+        "revision": "revision",
+        "runtime_dtype": torch.float32,
+        "load_format": "default",
+        "loader_inputs": {"model_config": {"hidden_size": 3}},
+        "tensor_parallel_size": 1,
+        "tensor_parallel_rank": 0,
+        "sequence_parallel_guard": {"sequence_parallel_size": 1, "backend": "none"},
+        "use_hsdp": False,
+        "enable_expert_parallel": False,
+        "quantization_config": None,
+        "cfg_parallel_size": 1,
+        "pipeline_parallel_size": 1,
+    }
+    kwargs.update(overrides)
+    return build_runtime_weight_cache_plan(pipeline, **kwargs)
+
+
+def _multiprocess_builder(cache_root: str, queue) -> None:
+    result = _build(_Pipeline(), Path(cache_root))
+    queue.put(
+        (
+            result.fallback_code,
+            result.plan.runtime_layout_key if result.plan is not None else None,
+            sorted(binding.file_path for binding in result.plan.bindings.values()) if result.plan is not None else [],
+        )
+    )
+
+
+def test_equivalent_process_layouts_reuse_one_entry(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    second = _build(_Pipeline(), tmp_path)
+
+    assert first.plan is not None
+    assert second.plan is not None
+    assert first.plan.runtime_layout_key == second.plan.runtime_layout_key
+    assert first.plan.post_load_complete
+    assert {binding.file_path for binding in first.plan.bindings.values()} == {
+        binding.file_path for binding in second.plan.bindings.values()
+    }
+    assert set(first.plan.bindings) == {
+        "transformer.persistent",
+        "transformer.proj.bias",
+        "transformer.proj.weight",
+    }
+
+
+def test_tp_coordinate_and_sp_implementation_guard_split_entries(tmp_path):
+    pipeline = _Pipeline()
+    tp0 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=0)
+    tp1 = _build(pipeline, tmp_path, tensor_parallel_size=2, tensor_parallel_rank=1)
+    sp2 = _build(
+        pipeline,
+        tmp_path,
+        sequence_parallel_guard={"sequence_parallel_size": 2, "backend": "ulysses"},
+    )
+
+    assert tp0.plan is not None and tp1.plan is not None and sp2.plan is not None
+    assert len({tp0.plan.runtime_layout_key, tp1.plan.runtime_layout_key, sp2.plan.runtime_layout_key}) == 3
+
+
+def test_final_runtime_content_is_authoritative(tmp_path):
+    first_pipeline = _Pipeline()
+    second_pipeline = _Pipeline()
+    with torch.no_grad():
+        second_pipeline.transformer.proj.weight[0, 0] += 1
+
+    first = _build(first_pipeline, tmp_path)
+    second = _build(second_pipeline, tmp_path)
+
+    assert first.plan is not None and second.plan is not None
+    assert first.plan.runtime_layout_key != second.plan.runtime_layout_key
+
+
+def test_expert_parallel_ownership_fails_closed(tmp_path):
+    result = _build(_Pipeline(), tmp_path, enable_expert_parallel=True)
+
+    assert result.plan is None
+    assert result.fallback_code == "unsupported_expert_parallel"
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_corrupt_published_entry_is_rebuilt_under_the_key_lock(tmp_path):
+    first = _build(_Pipeline(), tmp_path)
+    assert first.plan is not None
+    shard = Path(next(iter(first.plan.bindings.values())).file_path)
+    old_inode = shard.stat().st_ino
+    with shard.open("r+b") as handle:
+        handle.seek(-1, 2)
+        byte = handle.read(1)
+        handle.seek(-1, 2)
+        handle.write(bytes([byte[0] ^ 0xFF]))
+
+    second = _build(_Pipeline(), tmp_path)
+
+    assert second.plan is not None
+    assert second.plan.runtime_layout_key == first.plan.runtime_layout_key
+    assert Path(next(iter(second.plan.bindings.values())).file_path).stat().st_ino != old_inode
+
+
+def test_lock_timeout_keeps_the_ordinary_tensors(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    lock_path = tmp_path / ".locks" / f"{cache_key}.lock"
+
+    with lock_path.open("a+b") as lock_handle:
+        fcntl.flock(lock_handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+        result = _build(_Pipeline(), tmp_path, lock_timeout_seconds=0.01)
+
+    assert result.plan is None
+    assert result.fallback_code == "lock_timeout"
+
+
+def test_next_writer_removes_stale_temp_after_process_death(tmp_path):
+    initial = _build(_Pipeline(), tmp_path)
+    assert initial.plan is not None
+    cache_key = initial.plan.runtime_layout_key
+    entry_dir = Path(next(iter(initial.plan.bindings.values())).file_path).parent
+    shutil.rmtree(entry_dir)
+    stale_temp = tmp_path / ".tmp" / f"{cache_key}.stale-writer"
+    stale_temp.mkdir(parents=True)
+    (stale_temp / "partial").write_bytes(b"partial")
+
+    result = _build(_Pipeline(), tmp_path)
+
+    assert result.plan is not None
+    assert not stale_temp.exists()
+
+
+@pytest.mark.parametrize("case", ["alias", "external_alias", "noncontiguous"])
+def test_unsupported_tensor_layouts_fail_closed(tmp_path, case):
+    pipeline = _Pipeline()
+    if case == "alias":
+        pipeline.transformer.tied = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    elif case == "external_alias":
+        pipeline.encoder = nn.Module()
+        pipeline.encoder.shared = pipeline.transformer.proj.weight
+        expected_code = "unsupported_alias"
+    else:
+        pipeline.transformer.proj.weight = nn.Parameter(torch.arange(6, dtype=torch.float32).reshape(3, 2).t())
+        expected_code = "unsupported_tensor"
+
+    result = _build(pipeline, tmp_path)
+
+    assert result.plan is None
+    assert result.fallback_code == expected_code
+    assert list(tmp_path.rglob("*.safetensors")) == []
+
+
+def test_concurrent_processes_publish_one_immutable_entry(tmp_path):
+    context = mp.get_context("spawn")
+    queue = context.Queue()
+    processes = [context.Process(target=_multiprocess_builder, args=(str(tmp_path), queue)) for _ in range(2)]
+    for process in processes:
+        process.start()
+    results = [queue.get(timeout=30) for _ in processes]
+    for process in processes:
+        process.join(timeout=30)
+
+    assert all(process.exitcode == 0 for process in processes)
+    assert all(code is None for code, _, _ in results)
+    assert len({key for _, key, _ in results}) == 1
+    assert len({tuple(paths) for _, _, paths in results}) == 1
+    assert len(list((tmp_path / "v1").iterdir())) == 1

tests/diffusion/offloader/test_cuda_host_registration.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+
+from __future__ import annotations
+
+from types import SimpleNamespace
+
+import pytest
+
+import vllm_omni.diffusion.offloader.cuda_host_registration as registration_module
+from vllm_omni.diffusion.offloader.cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationBudgetError,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+    _coalesce_ranges,
+)
+
+pytestmark = [pytest.mark.diffusion, pytest.mark.cpu, pytest.mark.core_model]
+
+
+class _FakeStorage:
+    def __init__(self, pointer: int, size: int) -> None:
+        self._pointer = pointer
+        self._size = size
+
+    def data_ptr(self) -> int:
+        return self._pointer
+
+    def nbytes(self) -> int:
+        return self._size
+
+
+class _FakeTensor:
+    device = SimpleNamespace(type="cpu")
+
+    def __init__(self, pointer: int, size: int, *, pinned: bool = True) -> None:
+        self._storage = _FakeStorage(pointer, size)
+        self._pinned = pinned
+
+    def untyped_storage(self) -> _FakeStorage:
+        return self._storage
+
+    def numel(self) -> int:
+        return self._storage.nbytes()
+
+    def is_pinned(self) -> bool:
+        return self._pinned
+
+
+class _FakeRuntime:
+    def __init__(
+        self,
+        register_results: list[int | Exception],
+        unregister_results: list[int | Exception] | None = None,
+    ) -> None:
+        self._register_results = iter(register_results)
+        self._unregister_results = iter(unregister_results or [])
+        self.registered: list[tuple[int, int, int]] = []
+        self.unregistered: list[int] = []
+
+    def cudaHostRegister(self, pointer: int, size: int, flags: int) -> int:
+        self.registered.append((pointer, size, flags))
+        result = next(self._register_results)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    def cudaHostUnregister(self, pointer: int) -> int:
+        self.unregistered.append(pointer)
+        result = next(self._unregister_results, 0)
+        if isinstance(result, Exception):
+            raise result
+        return result
+
+    @staticmethod
+    def cudaGetErrorString(error: int) -> str:
+        return f"error-{error}"
+
+
+def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
+    assert _coalesce_ranges(
+        [(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
+        page_size=4096,
+    ) == (
+        registration_module._AddressRange(0x1000, 0x3000),
+        registration_module._AddressRange(0x9000, 0xA000),
+    )
+
+
+def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationBudgetError, match="exceeding"):
+        CudaHostRegistration.create(
+            {"weights": [_FakeTensor(0x1003, 4096)]},  # type: ignore[list-item]
+            max_bytes=4096,
+        )
+
+    assert runtime.registered == []
+
+
+def test_registration_rolls_back_partial_success(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="error-7"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_wraps_runtime_exception_and_rolls_back(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, RuntimeError("driver rejected mapping")])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationError, match="driver rejected mapping"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+    assert runtime.unregistered == [0x1000]
+
+
+def test_registration_fails_closed_when_rollback_fails(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 7], unregister_results=[9])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    with pytest.raises(CudaHostRegistrationCleanupError, match="rollback errors"):
+        CudaHostRegistration.create(
+            {
+                "first": [_FakeTensor(0x1003, 1)],  # type: ignore[list-item]
+                "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+            },
+            max_bytes=8192,
+        )
+
+
+def test_registration_closes_all_coalesced_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {
+            "first": [_FakeTensor(0x1003, 4096), _FakeTensor(0x2800, 1024)],  # type: ignore[list-item]
+            "second": [_FakeTensor(0x9003, 1)],  # type: ignore[list-item]
+        },
+        max_bytes=12288,
+    )
+
+    assert registration.total_bytes == 12288
+    assert registration.region_count == 2
+    assert registration.close() == []
+    assert runtime.unregistered == [0x9000, 0x1000]
+    assert registration.close() == []
+
+
+def test_registration_keeps_adjacent_file_mappings_separate(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+
+    registration = CudaHostRegistration.create(
+        {
+            "first.safetensors": [_FakeTensor(0x1000, 4096)],  # type: ignore[list-item]
+            "second.safetensors": [_FakeTensor(0x2000, 4096)],  # type: ignore[list-item]
+        },
+        max_bytes=8192,
+    )
+
+    assert registration.region_count == 2
+    assert runtime.registered == [
+        (0x1000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+        (0x2000, 4096, registration_module._CUDA_HOST_REGISTER_READ_ONLY),
+    ]
+    assert registration.close() == []
+
+
+def test_registration_close_retries_failed_ranges(monkeypatch: pytest.MonkeyPatch) -> None:
+    monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
+    runtime = _FakeRuntime([0], unregister_results=[9, 0])
+    monkeypatch.setattr(registration_module.torch.cuda, "cudart", lambda: runtime)
+    registration = CudaHostRegistration.create(
+        {"weights": [_FakeTensor(0x1003, 1)]},  # type: ignore[list-item]
+        max_bytes=4096,
+    )
+
+    assert registration.close() == ["cudaHostUnregister(0x1000) failed: error-9"]
+    assert registration.close() == []
+    assert runtime.unregistered == [0x1000, 0x1000]

tests/diffusion/offloader/test_distributed_layerwise_backend.py

@@ -473,6 +473,31 @@ class TestDistributedLayerwiseOffloadHook:
         assert staging[0][torch.float32].numel() == 12
         assert staging[1][torch.float32].numel() == 12
 
+    def test_registered_mmap_copies_directly_without_staging(self, patched_offload_runtime):
+        current_block = nn.Linear(2, 2, bias=False)
+        next_block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        next_block.weight.data.copy_(expected)
+        hook = DistributedLayerwiseOffloadHook(
+            next_block=next_block,
+            device=torch.device("cpu"),
+            dp_group=None,
+            dp_size=1,
+            rank=0,
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        hook.initialize_hook(current_block)
+        hook.registered_mmap = True
+
+        def fail_staging(_slot):
+            raise AssertionError("registered mmap must bypass host staging")
+
+        hook._stage_mmap_sources = fail_staging  # type: ignore[method-assign]
+        hook.prefetch_layer(slot=0, non_blocking=True)
+
+        assert torch.equal(next_block.weight, expected)
+
 
 class TestPinnedResidentLayerGroup:
     def test_load_offload_reuses_pinned_master_weights(self, patched_offload_runtime):
@@ -563,6 +588,24 @@ class TestPinnedResidentLayerGroup:
         group.offload()
         assert all(block.weight.numel() == 0 for block in blocks)
 
+    def test_registered_mmap_resident_layers_bypass_staging(self, patched_offload_runtime):
+        block = nn.Linear(2, 2, bias=False)
+        expected = torch.arange(4, dtype=torch.float32).view(2, 2)
+        block.weight.data.copy_(expected)
+        group = PinnedResidentLayerGroup(
+            [block],
+            device=torch.device("cpu"),
+            copy_stream=DummyStream(),
+            pin_memory=False,
+            rank_local_mmap=True,
+        )
+        group.registered_mmap = True
+        group._cpu_staging_buffers.clear()
+
+        group.load()
+
+        assert torch.equal(block.weight, expected)
+
     def test_shared_allgather_output_is_narrowed_to_current_block(
         self, dist_group, patched_offload_runtime, monkeypatch
     ):
@@ -692,6 +735,95 @@ class _MmapPostLoadPipeline(nn.Module):
 
 
 class TestMmapWeightLoading:
+    def test_runtime_cache_maps_final_weights_without_rerunning_post_load(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        class Transformer(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.proj = nn.Linear(2, 2, bias=False)
+                self.post_load_calls = 0
+
+            def post_load_weights(self):
+                self.post_load_calls += 1
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        old_parameter = pipeline.transformer.proj.weight
+        old_parameter.runtime_marker = "preserved"
+        weight_file = tmp_path / "runtime.safetensors"
+        save_file({"transformer.proj.weight": torch.full((2, 2), 9.0)}, str(weight_file))
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={
+                "transformer.proj.weight": TensorBinding(
+                    storage_key="transformer.proj.weight",
+                    file_path=str(weight_file),
+                )
+            },
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        backend._load_weights_from_runtime_cache(pipeline, modules, plan)
+
+        assert pipeline.transformer.post_load_calls == 0
+        assert pipeline.transformer.proj.weight is old_parameter
+        assert pipeline.transformer.proj.weight.runtime_marker == "preserved"
+        assert torch.equal(pipeline.transformer.proj.weight, torch.full((2, 2), 9.0))
+        backend._release_mmap_handles()
+
+    def test_runtime_cache_mapping_failure_leaves_ordinary_tensors_intact(
+        self,
+        tmp_path,
+        patched_offload_runtime,
+    ):
+        pipeline = nn.Module()
+        pipeline.transformer = nn.Linear(2, 2, bias=False)
+        old_parameter = pipeline.transformer.weight
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={
+                "transformer.weight": TensorBinding(
+                    storage_key="transformer.weight",
+                    file_path=str(tmp_path / "missing.safetensors"),
+                )
+            },
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        modules = SimpleNamespace(dits=[pipeline.transformer], dit_names=["transformer"])
+
+        with pytest.raises(Exception, match="missing.safetensors"):
+            backend._load_weights_from_runtime_cache(pipeline, modules, plan)
+
+        assert pipeline.transformer.weight is old_parameter
+
     def test_runs_model_post_load_hook(self, tmp_path, patched_offload_runtime):
         pipeline = _MmapPostLoadPipeline()
         weights = {name: torch.ones(param.shape, dtype=torch.bfloat16) for name, param in pipeline.named_parameters()}
@@ -715,7 +847,7 @@ class TestMmapWeightLoading:
             backing_kind="checkpoint_mmap",
             bindings={
                 name: TensorBinding(
-                    checkpoint_key=name,
+                    storage_key=name,
                     file_path=str(weight_file),
                 )
                 for name in weights
@@ -786,6 +918,99 @@ class TestMmapWeightLoading:
         assert all(hook.rank_local_mmap for group in backend._all_hook_groups for hook in group)
         backend.disable()
 
+    def test_runtime_cache_enable_reclaims_private_weights_before_hook_setup(
+        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) for _ in range(2)])
+
+        class Pipeline(nn.Module):
+            def __init__(self):
+                super().__init__()
+                self.transformer = Transformer()
+
+        pipeline = Pipeline()
+        weights = {name: torch.full_like(param, 4) for name, param in pipeline.named_parameters()}
+        weight_file = tmp_path / "runtime.safetensors"
+        save_file(weights, str(weight_file))
+        plan = HostWeightPlan(
+            backing_kind="runtime_cache",
+            bindings={name: TensorBinding(storage_key=name, file_path=str(weight_file)) for name in weights},
+            runtime_layout_key="runtime-layout",
+            post_load_complete=True,
+        )
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=False,
+                dlo_use_allgather=False,
+            ),
+            torch.device("cpu"),
+            host_weight_plan=plan,
+        )
+        cleanup_calls: list[None] = []
+        backend._cleanup_after_loading = lambda: cleanup_calls.append(None)  # type: ignore[method-assign]
+
+        backend.enable(pipeline)
+
+        assert backend._using_rank_local_mmap
+        # One reclamation immediately after remap, plus the existing final
+        # post-hook cleanup.
+        assert len(cleanup_calls) == 2
+        backend.disable()
+
+    def test_runtime_cache_registration_uses_budget_and_releases_before_mmap(
+        self,
+        monkeypatch: pytest.MonkeyPatch,
+        patched_offload_runtime,
+    ):
+        events: list[str] = []
+
+        class Registration:
+            total_bytes = 4096
+            region_count = 1
+
+            @staticmethod
+            def close() -> list[str]:
+                events.append("unregister")
+                return []
+
+        create_calls: list[tuple[dict[str, list[torch.Tensor]], int]] = []
+
+        def create(sources, *, max_bytes):
+            create_calls.append((sources, max_bytes))
+            return Registration()
+
+        monkeypatch.setattr(dist_backend_module.CudaHostRegistration, "create", staticmethod(create))
+        backend = DistributedLayerwiseOffloadBackend(
+            OffloadConfig(
+                strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+                pin_cpu_memory=True,
+                dlo_use_allgather=False,
+                dlo_runtime_cache_pin_limit_gib=1.5,
+            ),
+            torch.device("cuda"),
+        )
+        sources = {"runtime.safetensors": [torch.ones(1)]}
+        backend._runtime_cache_mapped_sources = sources
+
+        assert backend._try_register_runtime_cache_mmap()
+        assert create_calls == [(sources, int(1.5 * 1024**3))]
+
+        backend.enabled = True
+        backend._using_rank_local_mmap = True
+        backend._mmap_file_cache = {}
+        backend._release_mmap_handles = lambda: events.append("mmap")  # type: ignore[method-assign]
+        backend.disable()
+
+        assert events == ["unregister", "mmap"]
+
 
 class TestGetBlocksFromDit:
     def test_get_blocks_from_dit_single_block_attr(self):
@@ -1091,7 +1316,7 @@ class TestMmapValidation:
         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",
+            storage_key="blocks.0.weight",
             file_path=str(checkpoint_file),
         )
         assert result.plan.planned_source_prefixes == frozenset({"transformer."})
@@ -1381,7 +1606,7 @@ class TestMmapValidation:
 
         assert result.plan is not None
         assert result.plan.bindings["transformer.proj_in.weight"] == TensorBinding(
-            checkpoint_key="proj_in.weight",
+            storage_key="proj_in.weight",
             file_path=str(checkpoint_file),
         )
 

tests/diffusion/test_worker_wrapper_base.py

@@ -12,11 +12,13 @@ This module tests the WorkerWrapperBase implementation:
 - Dynamic worker class extension
 """
 
+from types import SimpleNamespace
 from typing import Any
 
 import pytest
 from pytest_mock import MockerFixture
 
+import vllm_omni.diffusion.worker.diffusion_worker as worker_module
 from vllm_omni.diffusion.worker.diffusion_worker import (
     CustomPipelineWorkerExtension,
     DiffusionWorker,
@@ -204,6 +206,29 @@ class TestWorkerWrapperBaseDelegation:
         wrapper.worker.shutdown.assert_called_once()
         assert result is None
 
+    def test_worker_shutdown_disables_offloader_before_distributed_teardown(
+        self,
+        mocker: MockerFixture,
+        monkeypatch: pytest.MonkeyPatch,
+    ):
+        worker = object.__new__(DiffusionWorker)
+        offload_backend = mocker.Mock()
+        kv_transfer_manager = mocker.Mock()
+        worker.model_runner = SimpleNamespace(
+            offload_backend=offload_backend,
+            kv_transfer_manager=kv_transfer_manager,
+        )
+        events: list[str] = []
+        offload_backend.disable.side_effect = lambda: events.append("offloader")
+        kv_transfer_manager.shutdown_prefetch.side_effect = lambda: events.append("kv_transfer")
+        monkeypatch.setattr(worker_module, "destroy_distributed_env", lambda: events.append("distributed"))
+
+        worker.shutdown()
+
+        offload_backend.disable.assert_called_once_with()
+        kv_transfer_manager.shutdown_prefetch.assert_called_once_with()
+        assert events == ["offloader", "kv_transfer", "distributed"]
+
 
 # -------------------------------------------------------------------------
 # Tests: execute_method

tests/entrypoints/test_async_omni_diffusion_config.py

@@ -328,6 +328,13 @@ def test_serve_cli_forwards_distributed_offload_residency():
             "--omni",
             "--enable-distributed-layerwise-offload",
             "--dlo-no-use-allgather",
+            "--dlo-enable-runtime-cache",
+            "--dlo-runtime-cache-dir",
+            "/var/cache/vllm-omni-dlo",
+            "--dlo-runtime-cache-lock-timeout",
+            "30",
+            "--dlo-runtime-cache-pin-limit-gib",
+            "80",
             "--dlo-resident-layers",
             "20",
         ]
@@ -339,9 +346,17 @@ def test_serve_cli_forwards_distributed_offload_residency():
 
     assert args.enable_distributed_layerwise_offload is True
     assert args.dlo_use_allgather is False
+    assert args.dlo_enable_runtime_cache is True
+    assert args.dlo_runtime_cache_dir == "/var/cache/vllm-omni-dlo"
+    assert args.dlo_runtime_cache_lock_timeout == 30
+    assert args.dlo_runtime_cache_pin_limit_gib == 80
     assert args.dlo_resident_layers == 20
     assert engine_args["enable_distributed_layerwise_offload"] is True
     assert engine_args["dlo_use_allgather"] is False
+    assert engine_args["dlo_enable_runtime_cache"] is True
+    assert engine_args["dlo_runtime_cache_dir"] == "/var/cache/vllm-omni-dlo"
+    assert engine_args["dlo_runtime_cache_lock_timeout"] == 30
+    assert engine_args["dlo_runtime_cache_pin_limit_gib"] == 80
     assert engine_args["dlo_resident_layers"] == 20
 
 

vllm_omni/config/omni_config.py

@@ -576,6 +576,10 @@ class _DiffusionConfigProjection:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_enable_runtime_cache: bool = False
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float = Field(default=600.0, gt=0)
+    dlo_runtime_cache_pin_limit_gib: float = Field(default=0.0, ge=0)
     dlo_resident_layers: int = Field(default=0, ge=0)
     pin_cpu_memory: bool = True
     diffusion_compile_granularity: Literal["regional", "full"] = "regional"
@@ -638,6 +642,13 @@ class _DiffusionConfigProjection:
         from vllm_omni.diffusion.diffusion_kv.config import parse_diffusion_kv_cache_mode
         from vllm_omni.quantization import build_quant_config
 
+        if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
+            raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
+        if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
+            raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
+
         if self.tf_model_config is None:
             self.tf_model_config = TransformerConfig()
         elif isinstance(self.tf_model_config, Mapping):

vllm_omni/config/stage_config.py

@@ -382,6 +382,10 @@ class StageDeployConfig:
 
     enable_distributed_layerwise_offload: bool | None = None
     dlo_use_allgather: bool | None = None
+    dlo_enable_runtime_cache: bool | None = None
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float | None = None
+    dlo_runtime_cache_pin_limit_gib: float | None = None
     dlo_resident_layers: int | None = None
     # Diffusion-specific debug and observability knobs.
     enable_diffusion_pipeline_profiler: bool | None = None

vllm_omni/diffusion/data.py

@@ -708,6 +708,18 @@ class OmniDiffusionConfig:
     # This avoids AllGather synchronization, while host memory follows the
     # loader's existing rank-local layout instead of adding a second DP shard.
     dlo_use_allgather: bool = True
+    # Build or join a node-local mmap cache after ordinary loading. This is
+    # used when direct checkpoint mmap is unavailable, or when direct CUDA
+    # registration requires a transform-complete final layout. The first
+    # version is no-AllGather only.
+    dlo_enable_runtime_cache: bool = False
+    # Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-runtime-weights.
+    dlo_runtime_cache_dir: str | None = None
+    # Maximum wait for the per-layout POSIX writer lock.
+    dlo_runtime_cache_lock_timeout: float = 600.0
+    # Maximum file-backed memory a worker may register with CUDA. Zero keeps
+    # the bounded mmap-to-pinned staging path.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     # Leading main-DiT blocks kept resident by distributed layerwise offload.
     dlo_resident_layers: int = 0
 
@@ -928,6 +940,16 @@ class OmniDiffusionConfig:
         )
 
     def __post_init__(self):
+        if not math.isfinite(self.dlo_runtime_cache_lock_timeout) or self.dlo_runtime_cache_lock_timeout <= 0:
+            raise ValueError("dlo_runtime_cache_lock_timeout must be positive")
+        if not math.isfinite(self.dlo_runtime_cache_pin_limit_gib) or self.dlo_runtime_cache_pin_limit_gib < 0:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib must be non-negative")
+        if self.dlo_enable_runtime_cache and not self.enable_distributed_layerwise_offload:
+            raise ValueError("dlo_enable_runtime_cache requires distributed layerwise offload")
+        if self.dlo_enable_runtime_cache and self.dlo_use_allgather:
+            raise ValueError("dlo_enable_runtime_cache currently requires dlo_use_allgather=False")
+        if self.dlo_runtime_cache_pin_limit_gib and not self.dlo_enable_runtime_cache:
+            raise ValueError("dlo_runtime_cache_pin_limit_gib requires dlo_enable_runtime_cache=True")
         if self.diffusion_compile_granularity not in {"regional", "full"}:
             raise ValueError(
                 "diffusion_compile_granularity must be 'regional' or 'full', "

vllm_omni/diffusion/model_loader/diffusers_loader.py

@@ -14,6 +14,7 @@ import torch
 from torch import nn
 from vllm.config.load import LoadConfig
 from vllm.logger import init_logger
+from vllm.model_executor.layers.linear import UnquantizedLinearMethod
 from vllm.model_executor.layers.quantization.base_config import QuantizeMethodBase
 from vllm.model_executor.model_loader.weight_utils import (
     download_safetensors_index_file_from_hf,
@@ -440,10 +441,16 @@ class DiffusersPipelineLoader:
                 _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
+                _prefer_registered_runtime_cache = (
+                    _dist_offload
+                    and not _use_ag
+                    and bool(getattr(self.od_config, "dlo_enable_runtime_cache", False))
+                    and float(getattr(self.od_config, "dlo_runtime_cache_pin_limit_gib", 0.0)) > 0
+                )
 
                 plan_result = None
                 weight_sources = self._get_weight_sources(model)
-                if _dist_offload:
+                if _dist_offload and not _prefer_registered_runtime_cache:
                     modules = ModuleDiscovery.discover(model)
                     plan_result = build_checkpoint_mmap_plan(
                         model,
@@ -455,6 +462,16 @@ class DiffusersPipelineLoader:
                         online_quantization=_has_online_quant,
                     )
                     self.host_weight_plan = plan_result.plan
+                elif _prefer_registered_runtime_cache:
+                    # Direct checkpoint bindings can carry deferred transforms
+                    # (for example grouped-QKV reordering), so registering the
+                    # raw checkpoint mapping would not make every recurrent H2D
+                    # source directly copyable. Materialize once, apply every
+                    # post-load mutation, then publish/register the final layout.
+                    logger.info(
+                        "DLO runtime-cache registration requested; using the ordinary loader "
+                        "once to publish a transform-complete final mmap layout"
+                    )
 
                 _skip_load = self.host_weight_plan is not None
 
@@ -491,7 +508,8 @@ class DiffusersPipelineLoader:
                         )
                     if _dist_offload and plan_result is not None:
                         logger.info(
-                            "DLO direct checkpoint mmap unavailable; using ordinary loader: %s",
+                            "DLO direct checkpoint mmap unavailable [%s]; using ordinary loader: %s",
+                            plan_result.fallback_code or "unknown",
                             plan_result.fallback_reason,
                         )
                     logger.debug("Loading weights on %s ...", load_device)
@@ -516,7 +534,107 @@ class DiffusersPipelineLoader:
                 logger.info("Quantization complete, offloaded model back to CPU")
 
         self._apply_skip_softmax_calibration(model)
-        return model.eval()
+        model.eval()
+        self._maybe_build_runtime_weight_cache(model, load_format)
+        return model
+
+    def _maybe_build_runtime_weight_cache(self, model: nn.Module, load_format: str) -> None:
+        """Publish final ordinary-loader DiT tensors for no-AllGather DLO."""
+        if self.host_weight_plan is not None:
+            return
+        if not getattr(self.od_config, "enable_distributed_layerwise_offload", False):
+            return
+        if not getattr(self.od_config, "dlo_enable_runtime_cache", False):
+            return
+        if getattr(self.od_config, "dlo_use_allgather", True):
+            logger.info("DLO runtime cache is no-AllGather only; leaving AllGather storage unchanged")
+            return
+
+        parallel_config = self.parallel_config
+        tp_size = int(getattr(parallel_config, "tensor_parallel_size", 1))
+        tp_rank = 0
+        if tp_size > 1:
+            try:
+                from vllm.distributed import get_tensor_model_parallel_rank
+
+                tp_rank = int(get_tensor_model_parallel_rank())
+            except Exception as exc:
+                logger.warning(
+                    "DLO runtime cache unavailable [parallel_rank_unavailable]: cannot resolve TP rank for TP=%d: %s",
+                    tp_size,
+                    exc,
+                )
+                return
+
+        modules = ModuleDiscovery.discover(model)
+        from vllm_omni.diffusion.model_loader.runtime_weight_cache import (
+            DEFAULT_LOCK_TIMEOUT_SECONDS,
+            DEFAULT_SHARD_SIZE_BYTES,
+            build_runtime_weight_cache_plan,
+        )
+
+        loader_inputs = {
+            "model_class_name": getattr(self.od_config, "model_class_name", None),
+            "model_config": getattr(self.od_config, "model_config", None),
+            "diffusers_load_kwargs": getattr(self.od_config, "diffusers_load_kwargs", None),
+            "custom_pipeline_args": getattr(self.od_config, "custom_pipeline_args", None),
+            "lora_path": getattr(self.od_config, "lora_path", None),
+            "lora_scale": getattr(self.od_config, "lora_scale", None),
+            "lora_backend": getattr(self.od_config, "lora_backend", None),
+        }
+        sp_guard = {
+            "sequence_parallel_size": int(getattr(parallel_config, "sequence_parallel_size", 1)),
+            "ulysses_degree": int(getattr(parallel_config, "ulysses_degree", 1)),
+            "ring_degree": int(getattr(parallel_config, "ring_degree", 1)),
+            "allgather_degree": int(getattr(parallel_config, "allgather_degree", 1)),
+            "ulysses_mode": getattr(parallel_config, "ulysses_mode", "strict"),
+        }
+        model_identity = getattr(self.od_config, "model", None)
+        runtime_quantization = self.quant_config
+        if runtime_quantization is None and any(
+            (quant_method := getattr(module, "quant_method", None)) is not None
+            and not isinstance(quant_method, UnquantizedLinearMethod)
+            for module in model.modules()
+        ):
+            runtime_quantization = "module_quant_method"
+        result = build_runtime_weight_cache_plan(
+            model,
+            dit_modules=tuple(zip(modules.dit_names, modules.dits)),
+            loader_type=type(self),
+            cache_root=getattr(self.od_config, "dlo_runtime_cache_dir", None),
+            lock_timeout_seconds=float(
+                getattr(self.od_config, "dlo_runtime_cache_lock_timeout", DEFAULT_LOCK_TIMEOUT_SECONDS)
+            ),
+            max_shard_bytes=DEFAULT_SHARD_SIZE_BYTES,
+            model_identity=str(model_identity) if model_identity is not None else None,
+            revision=getattr(self.od_config, "revision", None),
+            runtime_dtype=getattr(self.od_config, "dtype", None),
+            load_format=load_format,
+            loader_inputs=loader_inputs,
+            tensor_parallel_size=tp_size,
+            tensor_parallel_rank=tp_rank,
+            sequence_parallel_guard=sp_guard,
+            use_hsdp=bool(getattr(parallel_config, "use_hsdp", False)),
+            enable_expert_parallel=bool(getattr(parallel_config, "enable_expert_parallel", False)),
+            quantization_config=runtime_quantization,
+            cfg_parallel_size=int(getattr(parallel_config, "cfg_parallel_size", 1)),
+            pipeline_parallel_size=int(getattr(parallel_config, "pipeline_parallel_size", 1)),
+        )
+        if result.plan is not None:
+            self.host_weight_plan = result.plan
+            logger.info(
+                "DLO runtime cache selected (layout=%s, TP=%d/%d, SP=%d)",
+                result.plan.runtime_layout_key[:12] if result.plan.runtime_layout_key else "unknown",
+                tp_rank,
+                tp_size,
+                sp_guard["sequence_parallel_size"],
+            )
+        else:
+            logger.warning(
+                "DLO runtime cache unavailable [%s]; retaining ordinary host tensors: %s",
+                result.fallback_code or "unknown",
+                result.fallback_reason or "unspecified reason",
+            )
 
     @staticmethod
     def _request_offload_after_quant(model: nn.Module) -> int:

vllm_omni/diffusion/model_loader/host_weight_plan.py

@@ -28,7 +28,7 @@ TensorTransform = Callable[[torch.Tensor], torch.Tensor]
 class TensorBinding:
     """One runtime tensor backed by one safetensors entry."""
 
-    checkpoint_key: str
+    storage_key: str
     file_path: str
     transform: TensorTransform | None = None
 
@@ -40,6 +40,8 @@ class HostWeightPlan:
     backing_kind: str
     bindings: dict[str, TensorBinding]
     planned_source_prefixes: frozenset[str] = frozenset()
+    runtime_layout_key: str | None = None
+    post_load_complete: bool = False
 
 
 @dataclass(frozen=True)
@@ -48,6 +50,7 @@ class HostWeightPlanResult:
 
     plan: HostWeightPlan | None
     fallback_reason: str | None = None
+    fallback_code: str | None = None
 
 
 class _PlanIncompatibleError(RuntimeError):
@@ -258,11 +261,11 @@ def _validate_source_metadata(
         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:
+                if binding.storage_key not in available:
                     raise _PlanIncompatibleError(
-                        f"checkpoint key {binding.checkpoint_key!r} for {runtime_name!r} is missing from {file_path}"
+                        f"checkpoint key {binding.storage_key!r} for {runtime_name!r} is missing from {file_path}"
                     )
-                source = handle.get_slice(binding.checkpoint_key)
+                source = handle.get_slice(binding.storage_key)
                 source_shape = tuple(source.get_shape())
                 source_dtype = _SAFETENSORS_DTYPES.get(source.get_dtype())
                 if source_shape != tuple(target.shape):
@@ -287,11 +290,19 @@ def build_checkpoint_mmap_plan(
 ) -> 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")
+        return HostWeightPlanResult(
+            None,
+            f"TP={tensor_parallel_size} requires the ordinary loader",
+            "unsupported_tp",
+        )
     if use_hsdp:
-        return HostWeightPlanResult(None, "HSDP requires the ordinary loader")
+        return HostWeightPlanResult(None, "HSDP requires the ordinary loader", "unsupported_hsdp")
     if online_quantization:
-        return HostWeightPlanResult(None, "online quantization requires the ordinary loader")
+        return HostWeightPlanResult(
+            None,
+            "online quantization requires the ordinary loader",
+            "unsupported_online_quantization",
+        )
 
     remap_fn = getattr(type(pipeline), "_remap_ckpt_key", None)
     if not callable(remap_fn):
@@ -328,14 +339,14 @@ def build_checkpoint_mmap_plan(
                 )
             checkpoint_key, file_path = model_to_ckpt[runtime_name]
             bindings[runtime_name] = TensorBinding(
-                checkpoint_key=checkpoint_key,
+                storage_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(None, str(exc), "checkpoint_plan_incompatible")
 
     return HostWeightPlanResult(
         HostWeightPlan(

vllm_omni/diffusion/model_loader/runtime_weight_cache.py

@@ -0,0 +1,700 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Node-local, content-addressed runtime-weight cache for diffusion DLO."""
+
+from __future__ import annotations
+
+import contextlib
+import dataclasses
+import errno
+import fcntl
+import hashlib
+import inspect
+import json
+import math
+import os
+import shutil
+import time
+import uuid
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from itertools import chain
+from pathlib import Path
+from typing import Any
+
+import torch
+from safetensors import safe_open
+from safetensors.torch import save_file
+from torch import nn
+from vllm.logger import init_logger
+
+from vllm_omni.diffusion.model_loader.host_weight_plan import (
+    HostWeightPlan,
+    HostWeightPlanResult,
+    TensorBinding,
+)
+
+logger = init_logger(__name__)
+
+RUNTIME_CACHE_SCHEMA_VERSION = 1
+DEFAULT_RUNTIME_CACHE_ROOT = "~/.cache/vllm-omni/dlo-runtime-weights"
+DEFAULT_LOCK_TIMEOUT_SECONDS = 600.0
+DEFAULT_SHARD_SIZE_BYTES = 5 * 1024**3
+_HASH_CHUNK_BYTES = 64 * 1024**2
+_FILE_HASH_CHUNK_BYTES = 8 * 1024**2
+
+_SUPPORTED_DTYPES = {
+    torch.bool,
+    torch.uint8,
+    torch.int8,
+    torch.int16,
+    torch.int32,
+    torch.int64,
+    torch.float16,
+    torch.bfloat16,
+    torch.float32,
+    torch.float64,
+}
+for _dtype_name in ("float8_e4m3fn", "float8_e5m2"):
+    if (_dtype := getattr(torch, _dtype_name, None)) is not None:
+        _SUPPORTED_DTYPES.add(_dtype)
+
+
+@dataclass(frozen=True)
+class _RuntimeTensor:
+    name: str
+    tensor: torch.Tensor
+    kind: str
+
+
+class _RuntimeCacheError(RuntimeError):
+    def __init__(self, code: str, message: str):
+        super().__init__(message)
+        self.code = code
+
+
+def default_runtime_cache_root() -> str:
+    """Return the default root, intentionally independent of VLLM_CACHE_ROOT."""
+    return os.path.expanduser(DEFAULT_RUNTIME_CACHE_ROOT)
+
+
+def _persistent_buffer(module: nn.Module, local_name: str) -> bool:
+    parent_path, _, leaf_name = local_name.rpartition(".")
+    owner = module.get_submodule(parent_path)
+    return leaf_name not in owner._non_persistent_buffers_set
+
+
+def _named_parameters_with_duplicates(module: nn.Module):
+    try:
+        return module.named_parameters(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_parameters()
+
+
+def _named_buffers_with_duplicates(module: nn.Module):
+    try:
+        return module.named_buffers(remove_duplicate=False)
+    except TypeError:  # pragma: no cover - compatibility with old torch
+        return module.named_buffers()
+
+
+def _resolve_pipeline_tensor(pipeline: nn.Module, runtime_name: str) -> torch.Tensor | None:
+    parent_path, _, leaf_name = runtime_name.rpartition(".")
+    parent = pipeline.get_submodule(parent_path)
+    tensor = parent._parameters.get(leaf_name)
+    if tensor is None:
+        tensor = parent._buffers.get(leaf_name)
+    return tensor
+
+
+def _collect_runtime_tensors(
+    pipeline: nn.Module,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+) -> list[_RuntimeTensor]:
+    records: dict[str, _RuntimeTensor] = {}
+    for dit_name, dit_module in dit_modules:
+        for local_name, tensor in _named_parameters_with_duplicates(dit_module):
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = _RuntimeTensor(runtime_name, tensor, "parameter")
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise _RuntimeCacheError(
+                    "ambiguous_ownership",
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+        for local_name, tensor in _named_buffers_with_duplicates(dit_module):
+            if not _persistent_buffer(dit_module, local_name):
+                continue
+            runtime_name = f"{dit_name}.{local_name}"
+            candidate = _RuntimeTensor(runtime_name, tensor, "buffer")
+            existing = records.get(runtime_name)
+            if existing is not None and existing.tensor is not tensor:
+                raise _RuntimeCacheError(
+                    "ambiguous_ownership",
+                    f"multiple DiT tensors resolve to {runtime_name!r}",
+                )
+            records[runtime_name] = candidate
+
+    if not records:
+        raise _RuntimeCacheError(
+            "ambiguous_ownership",
+            "no DiT parameters or persistent buffers were discovered",
+        )
+
+    storage_owners: dict[tuple[int, int], str] = {}
+    for record in records.values():
+        tensor = record.tensor
+        try:
+            pipeline_tensor = _resolve_pipeline_tensor(pipeline, record.name)
+        except AttributeError as exc:
+            raise _RuntimeCacheError(
+                "ambiguous_ownership",
+                f"DiT tensor {record.name!r} is not owned by the pipeline",
+            ) from exc
+        if pipeline_tensor is not tensor:
+            raise _RuntimeCacheError(
+                "ambiguous_ownership",
+                f"DiT tensor {record.name!r} does not resolve to the discovered object",
+            )
+        if tensor.device.type != "cpu" or tensor.is_meta:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} must be a materialized CPU tensor, got {tensor.device}",
+            )
+        if tensor.layout != torch.strided:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} uses unsupported layout {tensor.layout}",
+            )
+        if not tensor.is_contiguous():
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is non-contiguous with stride {tensor.stride()}",
+            )
+        if tensor.dtype not in _SUPPORTED_DTYPES:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} uses unsupported dtype {tensor.dtype}",
+            )
+        if hasattr(tensor, "to_local"):
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is a distributed tensor",
+            )
+
+        tensor_nbytes = tensor.numel() * tensor.element_size()
+        storage = tensor.untyped_storage()
+        if tensor.storage_offset() != 0 or storage.nbytes() != tensor_nbytes:
+            raise _RuntimeCacheError(
+                "unsupported_tensor",
+                f"{record.name!r} is a view into a larger storage",
+            )
+        if tensor_nbytes:
+            storage_id = (storage.data_ptr(), storage.nbytes())
+            owner = storage_owners.setdefault(storage_id, record.name)
+            if owner != record.name:
+                raise _RuntimeCacheError(
+                    "unsupported_alias",
+                    f"{record.name!r} shares storage with {owner!r}",
+                )
+
+    # The ownership boundary is the complete pipeline, not only the discovered
+    # DiT traversal. Reject a DiT storage that is also registered by an encoder,
+    # VAE, resident component, or non-persistent buffer.
+    for pipeline_name, tensor in chain(
+        _named_parameters_with_duplicates(pipeline),
+        _named_buffers_with_duplicates(pipeline),
+    ):
+        if tensor.device.type != "cpu" or tensor.is_meta or tensor.numel() == 0:
+            continue
+        storage = tensor.untyped_storage()
+        owner = storage_owners.get((storage.data_ptr(), storage.nbytes()))
+        if owner is not None and owner != pipeline_name:
+            raise _RuntimeCacheError(
+                "unsupported_alias",
+                f"cached tensor {owner!r} shares storage with pipeline tensor {pipeline_name!r}",
+            )
+
+    return [records[name] for name in sorted(records)]
+
+
+def _json_normalize(value: Any) -> Any:
+    if value is None or isinstance(value, (bool, int, float, str)):
+        return value
+    if isinstance(value, torch.dtype):
+        return str(value)
+    if isinstance(value, os.PathLike):
+        return os.fspath(value)
+    if dataclasses.is_dataclass(value):
+        return _json_normalize(dataclasses.asdict(value))
+    model_dump = getattr(value, "model_dump", None)
+    if callable(model_dump):
+        try:
+            dumped = model_dump(mode="json")
+        except (TypeError, ValueError):
+            dumped = model_dump()
+        return _json_normalize(dumped)
+    if isinstance(value, Mapping):
+        return {str(key): _json_normalize(item) for key, item in sorted(value.items(), key=lambda item: str(item[0]))}
+    if isinstance(value, (set, frozenset)):
+        return sorted((_json_normalize(item) for item in value), key=lambda item: _canonical_json(item))
+    if isinstance(value, (list, tuple)):
+        return [_json_normalize(item) for item in value]
+    if inspect.isclass(value):
+        return f"{value.__module__}.{value.__qualname__}"
+    return {
+        "type": f"{type(value).__module__}.{type(value).__qualname__}",
+        "value": str(value),
+    }
+
+
+def _canonical_json(value: Any) -> bytes:
+    return json.dumps(
+        _json_normalize(value),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=True,
+    ).encode()
+
+
+def _implementation_fingerprint(
+    loader_type: type, pipeline: nn.Module, dit_modules: Sequence[tuple[str, nn.Module]]
+) -> str:
+    objects: list[Any] = [loader_type, type(pipeline)]
+    load_weights = getattr(type(pipeline), "load_weights", None)
+    if load_weights is not None:
+        objects.append(load_weights)
+    for _, dit_module in dit_modules:
+        objects.append(type(dit_module))
+        post_load = getattr(type(dit_module), "post_load_weights", None)
+        if post_load is not None:
+            objects.append(post_load)
+
+    digest = hashlib.sha256()
+    identities: set[str] = set()
+    for obj in objects:
+        identity = f"{getattr(obj, '__module__', '')}.{getattr(obj, '__qualname__', type(obj).__qualname__)}"
+        if identity in identities:
+            continue
+        identities.add(identity)
+        digest.update(identity.encode())
+        try:
+            digest.update(inspect.getsource(obj).encode())
+        except (OSError, TypeError):
+            pass
+    return digest.hexdigest()
+
+
+def _tensor_metadata(record: _RuntimeTensor) -> dict[str, Any]:
+    tensor = record.tensor
+    return {
+        "kind": record.kind,
+        "shape": list(tensor.shape),
+        "dtype": str(tensor.dtype),
+        "nbytes": tensor.numel() * tensor.element_size(),
+        "layout": "contiguous",
+    }
+
+
+def _update_hash_with_tensor(digest: Any, tensor: torch.Tensor) -> None:
+    byte_view = tensor.detach().reshape(-1).view(torch.uint8)
+    for offset in range(0, byte_view.numel(), _HASH_CHUNK_BYTES):
+        chunk = byte_view[offset : offset + _HASH_CHUNK_BYTES]
+        digest.update(memoryview(chunk.numpy()))
+
+
+def _runtime_content_digest(records: Sequence[_RuntimeTensor]) -> str:
+    digest = hashlib.sha256()
+    for record in records:
+        digest.update(_canonical_json({"name": record.name, **_tensor_metadata(record)}))
+        _update_hash_with_tensor(digest, record.tensor)
+    return digest.hexdigest()
+
+
+def _file_digest(path: Path) -> str:
+    digest = hashlib.sha256()
+    with path.open("rb") as handle:
+        while chunk := handle.read(_FILE_HASH_CHUNK_BYTES):
+            digest.update(chunk)
+    return digest.hexdigest()
+
+
+@contextlib.contextmanager
+def _exclusive_lock(path: Path, timeout_seconds: float):
+    path.parent.mkdir(parents=True, exist_ok=True)
+    deadline = time.monotonic() + timeout_seconds
+    with path.open("a+b") as handle:
+        while True:
+            try:
+                fcntl.flock(handle.fileno(), fcntl.LOCK_EX | fcntl.LOCK_NB)
+                break
+            except OSError as exc:
+                if exc.errno not in (errno.EACCES, errno.EAGAIN):
+                    raise _RuntimeCacheError("lock_failed", f"failed to lock {path}: {exc}") from exc
+                remaining = deadline - time.monotonic()
+                if remaining <= 0:
+                    raise _RuntimeCacheError(
+                        "lock_timeout",
+                        f"timed out after {timeout_seconds:g}s waiting for runtime-cache writer {path.name}",
+                    ) from exc
+                time.sleep(min(0.2, remaining))
+        try:
+            yield
+        finally:
+            fcntl.flock(handle.fileno(), fcntl.LOCK_UN)
+
+
+def _fsync_file(path: Path) -> None:
+    with path.open("rb") as handle:
+        os.fsync(handle.fileno())
+
+
+def _fsync_directory(path: Path) -> None:
+    fd = os.open(path, os.O_RDONLY | getattr(os, "O_DIRECTORY", 0))
+    try:
+        os.fsync(fd)
+    finally:
+        os.close(fd)
+
+
+def _split_shards(records: Sequence[_RuntimeTensor], max_shard_bytes: int) -> list[list[_RuntimeTensor]]:
+    shards: list[list[_RuntimeTensor]] = []
+    current: list[_RuntimeTensor] = []
+    current_bytes = 0
+    for record in records:
+        nbytes = record.tensor.numel() * record.tensor.element_size()
+        if current and current_bytes + nbytes > max_shard_bytes:
+            shards.append(current)
+            current = []
+            current_bytes = 0
+        current.append(record)
+        current_bytes += nbytes
+    if current:
+        shards.append(current)
+    return shards
+
+
+def _publish_entry(
+    *,
+    cache_root: Path,
+    entry_dir: Path,
+    cache_key: str,
+    layout_identity: dict[str, Any],
+    content_digest: str,
+    records: Sequence[_RuntimeTensor],
+    max_shard_bytes: int,
+) -> None:
+    total_nbytes = sum(record.tensor.numel() * record.tensor.element_size() for record in records)
+    free_bytes = shutil.disk_usage(cache_root).free
+    required_bytes = total_nbytes + max(64 * 1024**2, total_nbytes // 100)
+    if free_bytes < required_bytes:
+        raise _RuntimeCacheError(
+            "insufficient_disk",
+            f"runtime cache needs {required_bytes} bytes but only {free_bytes} are free under {cache_root}",
+        )
+
+    tmp_root = cache_root / ".tmp"
+    tmp_root.mkdir(parents=True, exist_ok=True)
+    temp_dir = tmp_root / f"{cache_key}.{os.getpid()}.{uuid.uuid4().hex}"
+    temp_dir.mkdir()
+    try:
+        shards = _split_shards(records, max_shard_bytes)
+        tensor_manifest: dict[str, dict[str, Any]] = {}
+        file_manifest: dict[str, dict[str, Any]] = {}
+        shard_count = len(shards)
+        for index, shard in enumerate(shards, start=1):
+            filename = f"model-{index:05d}-of-{shard_count:05d}.safetensors"
+            path = temp_dir / filename
+            save_file(
+                {record.name: record.tensor.detach() for record in shard},
+                str(path),
+                metadata={"format": "pt", "dlo_runtime_cache_schema": str(RUNTIME_CACHE_SCHEMA_VERSION)},
+            )
+            _fsync_file(path)
+            file_manifest[filename] = {
+                "size": path.stat().st_size,
+                "sha256": _file_digest(path),
+            }
+            for record in shard:
+                tensor_manifest[record.name] = {
+                    **_tensor_metadata(record),
+                    "file": filename,
+                    "storage_key": record.name,
+                }
+
+        manifest = {
+            "schema_version": RUNTIME_CACHE_SCHEMA_VERSION,
+            "cache_key": cache_key,
+            "layout_identity": layout_identity,
+            "runtime_content_sha256": content_digest,
+            "total_tensor_bytes": total_nbytes,
+            "created_at_unix_seconds": time.time(),
+            "writer_pid": os.getpid(),
+            "files": file_manifest,
+            "tensors": tensor_manifest,
+        }
+        manifest_path = temp_dir / "manifest.json"
+        with manifest_path.open("w", encoding="utf-8") as handle:
+            json.dump(manifest, handle, sort_keys=True, separators=(",", ":"))
+            handle.flush()
+            os.fsync(handle.fileno())
+        _fsync_directory(temp_dir)
+
+        invalid_dir: Path | None = None
+        if entry_dir.exists():
+            invalid_dir = tmp_root / f"{cache_key}.invalid.{os.getpid()}.{uuid.uuid4().hex}"
+            os.replace(entry_dir, invalid_dir)
+        os.replace(temp_dir, entry_dir)
+        _fsync_directory(entry_dir.parent)
+        if invalid_dir is not None:
+            shutil.rmtree(invalid_dir, ignore_errors=True)
+    except Exception:
+        shutil.rmtree(temp_dir, ignore_errors=True)
+        raise
+
+
+def _validate_entry(
+    *,
+    entry_dir: Path,
+    cache_key: str,
+    layout_identity: dict[str, Any],
+    content_digest: str,
+    expected_records: Sequence[_RuntimeTensor],
+) -> HostWeightPlan:
+    manifest_path = entry_dir / "manifest.json"
+    try:
+        with manifest_path.open(encoding="utf-8") as handle:
+            manifest = json.load(handle)
+    except (OSError, ValueError) as exc:
+        raise _RuntimeCacheError("entry_rejected", f"cannot read {manifest_path}: {exc}") from exc
+
+    if manifest.get("schema_version") != RUNTIME_CACHE_SCHEMA_VERSION:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache schema mismatch")
+    if manifest.get("cache_key") != cache_key:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache key mismatch")
+    if manifest.get("layout_identity") != layout_identity:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache layout identity mismatch")
+    if manifest.get("runtime_content_sha256") != content_digest:
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache content identity mismatch")
+
+    expected_metadata = {record.name: _tensor_metadata(record) for record in expected_records}
+    tensor_manifest = manifest.get("tensors")
+    file_manifest = manifest.get("files")
+    if not isinstance(tensor_manifest, dict) or not isinstance(file_manifest, dict):
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache manifest has invalid sections")
+    if set(tensor_manifest) != set(expected_metadata):
+        raise _RuntimeCacheError("entry_rejected", "runtime-cache tensor names differ from the loaded DiT")
+
+    for filename, metadata in file_manifest.items():
+        if Path(filename).name != filename or not isinstance(metadata, dict):
+            raise _RuntimeCacheError("entry_rejected", f"invalid cache filename {filename!r}")
+        path = entry_dir / filename
+        try:
+            stat = path.stat()
+        except OSError as exc:
+            raise _RuntimeCacheError("entry_rejected", f"missing cache shard {path}: {exc}") from exc
+        if stat.st_size != metadata.get("size") or _file_digest(path) != metadata.get("sha256"):
+            raise _RuntimeCacheError("entry_rejected", f"content digest mismatch for cache shard {path}")
+
+    bindings: dict[str, TensorBinding] = {}
+    mapped_records: list[_RuntimeTensor] = []
+    with contextlib.ExitStack() as stack:
+        handles = {
+            filename: stack.enter_context(safe_open(entry_dir / filename, framework="pt", device="cpu"))
+            for filename in file_manifest
+        }
+        expected_keys_by_file: dict[str, set[str]] = {filename: set() for filename in file_manifest}
+        for record in expected_records:
+            stored = tensor_manifest.get(record.name)
+            if not isinstance(stored, dict):
+                raise _RuntimeCacheError("entry_rejected", f"invalid tensor metadata for {record.name!r}")
+            filename = stored.get("file")
+            storage_key = stored.get("storage_key")
+            if filename not in handles or not isinstance(storage_key, str):
+                raise _RuntimeCacheError("entry_rejected", f"invalid storage binding for {record.name!r}")
+            actual_metadata = {key: stored.get(key) for key in expected_metadata[record.name]}
+            if actual_metadata != expected_metadata[record.name]:
+                raise _RuntimeCacheError("entry_rejected", f"metadata mismatch for {record.name!r}")
+            expected_keys_by_file[filename].add(storage_key)
+            try:
+                tensor = handles[filename].get_tensor(storage_key)
+            except Exception as exc:
+                raise _RuntimeCacheError(
+                    "entry_rejected",
+                    f"cannot map {record.name!r} from {filename}: {exc}",
+                ) from exc
+            if tuple(tensor.shape) != tuple(record.tensor.shape) or tensor.dtype != record.tensor.dtype:
+                raise _RuntimeCacheError("entry_rejected", f"mapped metadata mismatch for {record.name!r}")
+            mapped_records.append(_RuntimeTensor(record.name, tensor, record.kind))
+            bindings[record.name] = TensorBinding(storage_key=storage_key, file_path=str(entry_dir / filename))
+
+        for filename, expected_keys in expected_keys_by_file.items():
+            if set(handles[filename].keys()) != expected_keys:
+                raise _RuntimeCacheError("entry_rejected", f"unexpected tensor keys in cache shard {filename}")
+
+        if _runtime_content_digest(mapped_records) != content_digest:
+            raise _RuntimeCacheError("entry_rejected", "mapped runtime tensor content does not match the publisher")
+
+    return HostWeightPlan(
+        backing_kind="runtime_cache",
+        bindings=bindings,
+        runtime_layout_key=cache_key,
+        post_load_complete=True,
+    )
+
+
+def _remove_stale_temps(cache_root: Path, cache_key: str) -> None:
+    temp_root = cache_root / ".tmp"
+    if not temp_root.is_dir():
+        return
+    for path in temp_root.glob(f"{cache_key}.*"):
+        if path.is_dir():
+            shutil.rmtree(path, ignore_errors=True)
+        else:
+            path.unlink(missing_ok=True)
+
+
+def build_runtime_weight_cache_plan(
+    pipeline: nn.Module,
+    *,
+    dit_modules: Sequence[tuple[str, nn.Module]],
+    loader_type: type,
+    cache_root: str | os.PathLike[str] | None,
+    lock_timeout_seconds: float,
+    max_shard_bytes: int,
+    model_identity: str | None,
+    revision: str | None,
+    runtime_dtype: Any,
+    load_format: str,
+    loader_inputs: Mapping[str, Any],
+    tensor_parallel_size: int,
+    tensor_parallel_rank: int,
+    sequence_parallel_guard: Mapping[str, Any],
+    use_hsdp: bool,
+    enable_expert_parallel: bool,
+    quantization_config: Any,
+    cfg_parallel_size: int,
+    pipeline_parallel_size: int,
+) -> HostWeightPlanResult:
+    """Publish or join one immutable cache entry for final CPU DiT tensors."""
+    if load_format == "dummy":
+        return HostWeightPlanResult(
+            None,
+            "dummy/random weights do not have a reusable runtime identity",
+            "unsupported_load_format",
+        )
+    if use_hsdp:
+        return HostWeightPlanResult(None, "HSDP/DTensor layouts are not runtime-cache compatible", "unsupported_hsdp")
+    if enable_expert_parallel:
+        return HostWeightPlanResult(
+            None,
+            "expert-parallel weight ownership is outside runtime-cache v1",
+            "unsupported_expert_parallel",
+        )
+    if quantization_config is not None:
+        return HostWeightPlanResult(
+            None,
+            "quantized runtime layouts are not supported by the first runtime-cache version",
+            "unsupported_quantization",
+        )
+    if cfg_parallel_size != 1:
+        return HostWeightPlanResult(None, "CFG-parallel component ownership is not proven", "unsupported_cfg_parallel")
+    if pipeline_parallel_size != 1:
+        return HostWeightPlanResult(None, "pipeline-parallel ownership is outside runtime-cache v1", "unsupported_pp")
+    if not math.isfinite(lock_timeout_seconds) or lock_timeout_seconds <= 0:
+        return HostWeightPlanResult(None, "runtime-cache lock timeout must be positive", "invalid_configuration")
+    if max_shard_bytes <= 0:
+        return HostWeightPlanResult(None, "runtime-cache shard size must be positive", "invalid_configuration")
+
+    try:
+        records = _collect_runtime_tensors(pipeline, dit_modules)
+        content_digest = _runtime_content_digest(records)
+        layout_identity = {
+            "schema_version": RUNTIME_CACHE_SCHEMA_VERSION,
+            "model": model_identity,
+            "revision": revision,
+            "runtime_dtype": str(runtime_dtype),
+            "load_format": load_format,
+            "loader_inputs": _json_normalize(loader_inputs),
+            "loader_implementation_sha256": _implementation_fingerprint(loader_type, pipeline, dit_modules),
+            "components": [name for name, _ in dit_modules],
+            "tensor_parallel_size": tensor_parallel_size,
+            "tensor_parallel_rank": tensor_parallel_rank,
+            "sequence_parallel_guard": _json_normalize(sequence_parallel_guard),
+            "runtime_content_sha256": content_digest,
+        }
+        cache_key = hashlib.sha256(_canonical_json(layout_identity)).hexdigest()
+        root = Path(cache_root or default_runtime_cache_root()).expanduser().resolve()
+        entries_root = root / f"v{RUNTIME_CACHE_SCHEMA_VERSION}"
+        entry_dir = entries_root / cache_key
+        entries_root.mkdir(parents=True, exist_ok=True)
+
+        if entry_dir.is_dir():
+            try:
+                plan = _validate_entry(
+                    entry_dir=entry_dir,
+                    cache_key=cache_key,
+                    layout_identity=layout_identity,
+                    content_digest=content_digest,
+                    expected_records=records,
+                )
+                logger.info("Reusing DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
+                return HostWeightPlanResult(plan)
+            except _RuntimeCacheError as exc:
+                logger.warning("DLO runtime cache entry %s will be rebuilt: %s", cache_key[:12], exc)
+
+        lock_path = root / ".locks" / f"{cache_key}.lock"
+        with _exclusive_lock(lock_path, lock_timeout_seconds):
+            _remove_stale_temps(root, cache_key)
+            if entry_dir.is_dir():
+                try:
+                    plan = _validate_entry(
+                        entry_dir=entry_dir,
+                        cache_key=cache_key,
+                        layout_identity=layout_identity,
+                        content_digest=content_digest,
+                        expected_records=records,
+                    )
+                    logger.info("Joined DLO runtime cache entry %s from %s", cache_key[:12], entry_dir)
+                    return HostWeightPlanResult(plan)
+                except _RuntimeCacheError:
+                    pass
+
+            _publish_entry(
+                cache_root=root,
+                entry_dir=entry_dir,
+                cache_key=cache_key,
+                layout_identity=layout_identity,
+                content_digest=content_digest,
+                records=records,
+                max_shard_bytes=max_shard_bytes,
+            )
+            plan = _validate_entry(
+                entry_dir=entry_dir,
+                cache_key=cache_key,
+                layout_identity=layout_identity,
+                content_digest=content_digest,
+                expected_records=records,
+            )
+            logger.info(
+                "Published DLO runtime cache entry %s with %d tensors at %s",
+                cache_key[:12],
+                len(records),
+                entry_dir,
+            )
+            return HostWeightPlanResult(plan)
+    except _RuntimeCacheError as exc:
+        return HostWeightPlanResult(None, str(exc), exc.code)
+    except Exception as exc:
+        return HostWeightPlanResult(None, f"runtime-cache operation failed: {exc}", "cache_operation_failed")
+
+
+__all__ = [
+    "DEFAULT_LOCK_TIMEOUT_SECONDS",
+    "DEFAULT_RUNTIME_CACHE_ROOT",
+    "DEFAULT_SHARD_SIZE_BYTES",
+    "RUNTIME_CACHE_SCHEMA_VERSION",
+    "build_runtime_weight_cache_plan",
+    "default_runtime_cache_root",
+]

vllm_omni/diffusion/offloader/__init__.py

@@ -71,8 +71,9 @@ def get_offload_backend(
     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.
+        host_weight_plan: Exact loader-produced backing plan. Checkpoint plans
+            may have skipped ordinary materialization; runtime-cache plans
+            describe already-valid ordinary tensors that can be remapped.
 
     Returns:
         OffloadBackend instance or None if offloading disabled
@@ -85,10 +86,14 @@ def get_offload_backend(
     # Extract and validate configuration
     config = OffloadConfig.from_od_config(od_config)
 
+    plan_requires_backend = host_weight_plan is not None and not host_weight_plan.post_load_complete
     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"
-        )
+        if plan_requires_backend:
+            raise RuntimeError(
+                "A loader-owned DLO host-weight plan was produced, but distributed layerwise offload is not selected"
+            )
+        logger.warning("Ignoring optional runtime-cache plan because distributed layerwise offload is not selected")
+        host_weight_plan = None
 
     # Return None if no offloading requested
     if config.strategy == OffloadStrategy.NONE:
@@ -96,7 +101,7 @@ def get_offload_backend(
 
     # 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:
+        if plan_requires_backend:
             raise RuntimeError(
                 "The loader skipped ordinary weight materialization for DLO, "
                 "but this platform cannot create the required offload backend"
@@ -112,7 +117,7 @@ def get_offload_backend(
         try:
             device = current_omni_platform.get_torch_device()
         except (NotImplementedError, AttributeError) as exc:
-            if host_weight_plan is not None:
+            if plan_requires_backend:
                 raise RuntimeError(
                     "The loader skipped ordinary weight materialization for DLO, "
                     "but the target device could not be resolved"

vllm_omni/diffusion/offloader/base.py

@@ -31,6 +31,9 @@ class OffloadConfig:
     # 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
+    # Per-worker budget for registering shared runtime-cache mappings with CUDA.
+    # Zero retains bounded host staging.
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
 
     @classmethod
     def from_od_config(cls, od_config: OmniDiffusionConfig) -> "OffloadConfig":
@@ -101,6 +104,7 @@ class OffloadConfig:
         # requirements (concurrent requests, dummy run skip).
         dlo_use_allgather = getattr(od_config, "dlo_use_allgather", True)
         dlo_resident_layers = int(getattr(od_config, "dlo_resident_layers", 0))
+        dlo_runtime_cache_pin_limit_gib = float(getattr(od_config, "dlo_runtime_cache_pin_limit_gib", 0.0))
         if dlo_resident_layers < 0:
             raise ValueError(f"dlo_resident_layers must be >= 0, got {dlo_resident_layers}")
         if dlo_resident_layers and dlo_use_allgather:
@@ -136,6 +140,7 @@ class OffloadConfig:
             dp_size=dp_size,
             dlo_use_allgather=dlo_use_allgather,
             dlo_resident_layers=dlo_resident_layers,
+            dlo_runtime_cache_pin_limit_gib=dlo_runtime_cache_pin_limit_gib,
         )
 
 

vllm_omni/diffusion/offloader/cuda_host_registration.py

@@ -0,0 +1,205 @@
+# SPDX-License-Identifier: Apache-2.0
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
+"""Bounded CUDA registration for existing read-only host mappings."""
+
+from __future__ import annotations
+
+import mmap
+from collections.abc import Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any
+
+import torch
+
+_CUDA_HOST_REGISTER_READ_ONLY = 0x08
+
+
+class CudaHostRegistrationError(RuntimeError):
+    """Raised when an existing host range cannot be registered safely."""
+
+
+class CudaHostRegistrationBudgetError(CudaHostRegistrationError):
+    """Raised before mutation when the complete mapping exceeds its budget."""
+
+
+class CudaHostRegistrationCleanupError(CudaHostRegistrationError):
+    """Raised when partial registration could not be rolled back safely."""
+
+
+@dataclass(frozen=True)
+class _AddressRange:
+    start: int
+    end: int
+
+    @property
+    def size(self) -> int:
+        return self.end - self.start
+
+
+def _coalesce_ranges(
+    ranges: Sequence[tuple[int, int]],
+    page_size: int = mmap.PAGESIZE,
+) -> tuple[_AddressRange, ...]:
+    """Page-align and merge overlapping ranges from one backing mapping."""
+    if page_size <= 0:
+        raise ValueError(f"page_size must be positive, got {page_size}")
+
+    aligned: list[_AddressRange] = []
+    for start, size in ranges:
+        if start <= 0 or size < 0:
+            raise ValueError(f"invalid host range start={start}, size={size}")
+        if size == 0:
+            continue
+        aligned_start = start - start % page_size
+        end = start + size
+        aligned_end = ((end + page_size - 1) // page_size) * page_size
+        aligned.append(_AddressRange(aligned_start, aligned_end))
+
+    merged: list[_AddressRange] = []
+    for region in sorted(aligned, key=lambda item: (item.start, item.end)):
+        if merged and region.start <= merged[-1].end:
+            previous = merged[-1]
+            merged[-1] = _AddressRange(previous.start, max(previous.end, region.end))
+        else:
+            merged.append(region)
+    return tuple(merged)
+
+
+def _tensor_regions(
+    sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+) -> tuple[_AddressRange, ...]:
+    """Resolve storage spans without merging unrelated file mappings."""
+    regions: list[_AddressRange] = []
+    for mapping_name, tensors in sources_by_mapping.items():
+        ranges: list[tuple[int, int]] = []
+        for tensor in tensors:
+            if tensor.device.type != "cpu":
+                raise CudaHostRegistrationError(
+                    f"CUDA host registration requires CPU storage, but {mapping_name!r} contains {tensor.device}"
+                )
+            storage = tensor.untyped_storage()
+            ranges.append((storage.data_ptr(), storage.nbytes()))
+        regions.extend(_coalesce_ranges(ranges))
+    return tuple(regions)
+
+
+def _error_message(runtime: Any, error: Any) -> str:
+    try:
+        message = runtime.cudaGetErrorString(error)
+    except Exception:
+        return str(error)
+    if isinstance(message, bytes):
+        return message.decode(errors="replace")
+    return str(message)
+
+
+class CudaHostRegistration:
+    """Own CUDA registrations for already-existing file-backed CPU tensors."""
+
+    def __init__(
+        self,
+        runtime: Any,
+        regions: tuple[_AddressRange, ...],
+    ) -> None:
+        self._runtime = runtime
+        self._regions = regions
+        self._closed = False
+
+    @classmethod
+    def create(
+        cls,
+        sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
+        *,
+        max_bytes: int,
+    ) -> CudaHostRegistration:
+        if max_bytes <= 0:
+            raise CudaHostRegistrationBudgetError("CUDA host-registration budget is disabled")
+        regions = _tensor_regions(sources_by_mapping)
+        total_bytes = sum(region.size for region in regions)
+        if total_bytes > max_bytes:
+            raise CudaHostRegistrationBudgetError(
+                f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
+            )
+        if not regions:
+            raise CudaHostRegistrationError("no non-empty host ranges were available for registration")
+        if not torch.cuda.is_available():
+            raise CudaHostRegistrationError("CUDA is not available")
+
+        try:
+            runtime = torch.cuda.cudart()
+        except Exception as exc:
+            raise CudaHostRegistrationError(f"cannot access the CUDA runtime: {exc}") from exc
+
+        registered: list[_AddressRange] = []
+        try:
+            for region in regions:
+                error = runtime.cudaHostRegister(
+                    region.start,
+                    region.size,
+                    _CUDA_HOST_REGISTER_READ_ONLY,
+                )
+                if int(error) != 0:
+                    raise CudaHostRegistrationError(
+                        "cudaHostRegister(read-only) failed for "
+                        f"[{region.start:#x}, {region.end:#x}): {_error_message(runtime, error)}"
+                    )
+                registered.append(region)
+
+            unpinned = [
+                mapping_name
+                for mapping_name, tensors in sources_by_mapping.items()
+                if any(tensor.numel() and not tensor.is_pinned() for tensor in tensors)
+            ]
+            if unpinned:
+                raise CudaHostRegistrationError(
+                    f"CUDA registration succeeded but PyTorch did not recognize pinned storage for {unpinned[:3]}"
+                )
+        except Exception as exc:
+            rollback_errors: list[str] = []
+            for region in reversed(registered):
+                try:
+                    error = runtime.cudaHostUnregister(region.start)
+                    if int(error) != 0:
+                        rollback_errors.append(
+                            f"cudaHostUnregister({region.start:#x}) failed: {_error_message(runtime, error)}"
+                        )
+                except Exception as rollback_exc:
+                    rollback_errors.append(f"cudaHostUnregister({region.start:#x}) raised: {rollback_exc}")
+            if rollback_errors:
+                raise CudaHostRegistrationCleanupError(
+                    f"CUDA host registration failed ({exc}); rollback errors: {rollback_errors[:3]}"
+                ) from exc
+            if isinstance(exc, CudaHostRegistrationError):
+                raise
+            raise CudaHostRegistrationError(f"CUDA host registration raised: {exc}") from exc
+
+        return cls(runtime, regions)
+
+    @property
+    def total_bytes(self) -> int:
+        return sum(region.size for region in self._regions)
+
+    @property
+    def region_count(self) -> int:
+        return len(self._regions)
+
+    def close(self) -> list[str]:
+        """Unregister every range, returning errors after best-effort cleanup."""
+        if self._closed:
+            return []
+        errors: list[str] = []
+        failed: list[_AddressRange] = []
+        for region in reversed(self._regions):
+            try:
+                error = self._runtime.cudaHostUnregister(region.start)
+                if int(error) != 0:
+                    errors.append(
+                        f"cudaHostUnregister({region.start:#x}) failed: {_error_message(self._runtime, error)}"
+                    )
+                    failed.append(region)
+            except Exception as exc:
+                errors.append(f"cudaHostUnregister({region.start:#x}) raised: {exc}")
+                failed.append(region)
+        self._regions = tuple(reversed(failed))
+        self._closed = not self._regions
+        return errors

vllm_omni/diffusion/offloader/distributed_layerwise_backend.py

@@ -18,6 +18,7 @@ This module implements the RFC-1 "Distributed Layerwise Offload" mechanism that:
 
 from __future__ import annotations
 
+import time
 from itertools import chain
 from typing import Any
 
@@ -34,6 +35,11 @@ from vllm_omni.platforms import current_omni_platform
 
 from .base import OffloadBackend, OffloadConfig
 from .block_discovery import get_blocks_from_dit
+from .cuda_host_registration import (
+    CudaHostRegistration,
+    CudaHostRegistrationCleanupError,
+    CudaHostRegistrationError,
+)
 from .module_collector import ModuleDiscovery
 from .offload_plan import (
     OffloadPlan,
@@ -93,6 +99,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         self.rank = rank
         self.pin_memory = pin_memory
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.tensor_transforms = tensor_transforms or {}
 
         self.copy_stream = copy_stream or current_omni_platform.Stream()
@@ -444,6 +451,26 @@ class DistributedLayerwiseOffloadHook(ModelHook):
     #  Prefetch: H2D + AllGather (overlapped on dedicated streams)      #
     # ------------------------------------------------------------------ #
 
+    @staticmethod
+    def _resolve_mmap_source(
+        source_info: dict[str, Any],
+        meta: dict[str, Any],
+        dtype: torch.dtype,
+    ) -> torch.Tensor:
+        source = source_info["tensor"]
+        transform = source_info["transform"]
+        if callable(transform):
+            source = transform(source)
+        if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
+            raise ValueError(
+                "mmap weight transform changed tensor layout for "
+                f"{source_info['name']!r}: expected dtype={dtype}, "
+                f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
+                f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
+                f"stride={source.stride()}"
+            )
+        return source
+
     @staticmethod
     def _pack_mmap_sources(
         cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
@@ -456,18 +483,7 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             destination = slot_buffers[dtype][:total_numel]
             sources = cpu_sources[dtype]
             for source_info, meta in zip(sources, metas, strict=True):
-                source = source_info["tensor"]
-                transform = source_info["transform"]
-                if callable(transform):
-                    source = transform(source)
-                if source.dtype != dtype or source.shape != meta["shape"] or source.stride() != meta["stride"]:
-                    raise ValueError(
-                        "mmap weight transform changed tensor layout for "
-                        f"{source_info['name']!r}: expected dtype={dtype}, "
-                        f"shape={tuple(meta['shape'])}, stride={meta['stride']}; "
-                        f"got dtype={source.dtype}, shape={tuple(source.shape)}, "
-                        f"stride={source.stride()}"
-                    )
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
                 start = meta["offset"]
                 physical_storage = destination[start : start + meta["numel"]]
                 if source.is_contiguous():
@@ -481,6 +497,32 @@ class DistributedLayerwiseOffloadHook(ModelHook):
             staged[dtype] = destination
         return staged
 
+    @staticmethod
+    def _copy_mmap_sources_to_device(
+        cpu_sources: dict[torch.dtype, list[dict[str, Any]]],
+        metadata: dict[torch.dtype, list[dict[str, Any]]],
+        gpu_buffers: dict[torch.dtype, torch.Tensor],
+        *,
+        non_blocking: bool,
+    ) -> None:
+        """Copy registered source views directly into flattened device buffers."""
+        for dtype, metas in metadata.items():
+            destination = gpu_buffers[dtype]
+            sources = cpu_sources[dtype]
+            for source_info, meta in zip(sources, metas, strict=True):
+                source = DistributedLayerwiseOffloadHook._resolve_mmap_source(source_info, meta, dtype)
+                start = meta["offset"]
+                physical_storage = destination[start : start + meta["numel"]]
+                async_copy = non_blocking and source.is_pinned()
+                if source.is_contiguous():
+                    physical_storage.copy_(source.flatten(), non_blocking=async_copy)
+                else:
+                    torch.as_strided(
+                        physical_storage,
+                        size=source.shape,
+                        stride=source.stride(),
+                    ).copy_(source, non_blocking=async_copy)
+
     def _stage_mmap_sources(self, slot: int) -> dict[torch.dtype, torch.Tensor]:
         """Pack this block's mmap views into one bounded host staging slot."""
         previous_copy = self.cpu_staging_events[slot]
@@ -514,17 +556,27 @@ class DistributedLayerwiseOffloadHook(ModelHook):
         assert gpu_weights is not None, f"gpu_buffers[{slot}] not allocated"
 
         if self.dp_size <= 1 or self.dp_group is None:
-            cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
-            with current_omni_platform.stream(self.copy_stream):
-                for dtype, cpu_shard in cpu_weights.items():
-                    gw = gpu_weights[dtype]
-                    async_copy = non_blocking and cpu_shard.is_pinned()
-                    gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
-                evt.record(self.copy_stream)
-            if self.rank_local_mmap:
-                # The CPU slot may be overwritten only after this H2D copy has
-                # finished.  The shared event protects reuse by another hook.
-                self.cpu_staging_events[slot] = evt
+            if self.rank_local_mmap and self.registered_mmap:
+                with current_omni_platform.stream(self.copy_stream):
+                    self._copy_mmap_sources_to_device(
+                        self.cpu_sources,
+                        self.metadata,
+                        gpu_weights,
+                        non_blocking=non_blocking,
+                    )
+                    evt.record(self.copy_stream)
+            else:
+                cpu_weights = self._stage_mmap_sources(slot) if self.rank_local_mmap else self.cpu_shards
+                with current_omni_platform.stream(self.copy_stream):
+                    for dtype, cpu_shard in cpu_weights.items():
+                        gw = gpu_weights[dtype]
+                        async_copy = non_blocking and cpu_shard.is_pinned()
+                        gw[: cpu_shard.numel()].copy_(cpu_shard, non_blocking=async_copy)
+                    evt.record(self.copy_stream)
+                if self.rank_local_mmap:
+                    # The CPU slot may be overwritten only after this H2D copy has
+                    # finished.  The shared event protects reuse by another hook.
+                    self.cpu_staging_events[slot] = evt
         else:
             gpu_shards: dict[torch.dtype, torch.Tensor] = {}
             shard_bufs = self.gpu_shard_buffers[slot]
@@ -734,6 +786,7 @@ class PinnedResidentLayerGroup:
         self.copy_stream = copy_stream
         self.loaded = False
         self.rank_local_mmap = rank_local_mmap
+        self.registered_mmap = False
         self.pin_memory = pin_memory
         self._states: list[dict[str, Any]] = []
         self._gpu_buffers: list[dict[torch.dtype, torch.Tensor]] = []
@@ -804,6 +857,14 @@ class PinnedResidentLayerGroup:
         ready = current_omni_platform.Event()
         with current_omni_platform.stream(self.copy_stream):
             for index, (state, block_buffers) in enumerate(zip(self._states, gpu_buffers)):
+                if self.rank_local_mmap and self.registered_mmap:
+                    DistributedLayerwiseOffloadHook._copy_mmap_sources_to_device(
+                        state["cpu_sources"],
+                        state["metadata"],
+                        block_buffers,
+                        non_blocking=True,
+                    )
+                    continue
                 if self.rank_local_mmap:
                     slot = index % 2
                     previous_copy = self._cpu_staging_events[slot]
@@ -901,8 +962,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._resident_layer_group: PinnedResidentLayerGroup | None = None
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.host_weight_plan = host_weight_plan
         self._mmap_transforms_by_tensor_id: dict[int, Any] = {}
+        self._runtime_cache_mapped_sources: dict[str, list[torch.Tensor]] = {}
+        self._cuda_host_registration: CudaHostRegistration | None = None
 
     def load_resident_layers(self) -> None:
         """Load the model-declared leading blocks for the denoise stage."""
@@ -1023,7 +1087,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                     framework="pt",
                     device="cpu",
                 )
-            tensor = file_cache[binding.file_path].get_tensor(binding.checkpoint_key)
+            tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
             if is_parameter:
                 replacement = torch.nn.Parameter(tensor, requires_grad=target.requires_grad)
                 parent._parameters[leaf_name] = replacement
@@ -1103,6 +1167,154 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
                 f"(first 5: {remaining_meta[:5]})."
             )
 
+    def _load_weights_from_runtime_cache(
+        self,
+        pipeline: nn.Module,
+        modules,
+        plan: HostWeightPlan,
+    ) -> None:
+        """Replace final ordinary-loader tensors with validated cache views."""
+        from safetensors import safe_open
+
+        if not plan.post_load_complete:
+            raise RuntimeError("A runtime-cache plan must describe fully post-processed weights")
+        if any(binding.transform is not None for binding in plan.bindings.values()):
+            raise RuntimeError("Runtime-cache bindings cannot carry deferred checkpoint transforms")
+
+        required_names: set[str] = set()
+        for dit_name, dit_module in zip(modules.dit_names, modules.dits):
+            required_names.update(f"{dit_name}.{name}" for name, _ in dit_module.named_parameters())
+            for name, _ in dit_module.named_buffers():
+                parent_path, _, leaf_name = name.rpartition(".")
+                owner = dit_module.get_submodule(parent_path)
+                if leaf_name not in owner._non_persistent_buffers_set:
+                    required_names.add(f"{dit_name}.{name}")
+        if required_names != set(plan.bindings):
+            missing = sorted(required_names - set(plan.bindings))
+            unexpected = sorted(set(plan.bindings) - required_names)
+            raise RuntimeError(
+                "Runtime-cache plan coverage changed after loader validation: "
+                f"missing={missing[:5]}, unexpected={unexpected[:5]}"
+            )
+
+        file_cache: dict[str, Any] = {}
+        replacements: list[tuple[torch.Tensor, torch.Tensor]] = []
+        mapped_sources: dict[str, list[torch.Tensor]] = {}
+        try:
+            # Validate and map every tensor before mutating the pipeline. A
+            # missing/corrupt shard therefore leaves the ordinary tensors intact.
+            for runtime_name, binding in plan.bindings.items():
+                parent_path, _, leaf_name = runtime_name.rpartition(".")
+                try:
+                    parent = pipeline.get_submodule(parent_path)
+                except AttributeError as exc:
+                    raise RuntimeError(f"Runtime-cache target module {parent_path!r} no longer exists") from exc
+
+                target = parent._parameters.get(leaf_name)
+                if target is None:
+                    target = parent._buffers.get(leaf_name)
+                if target is None:
+                    raise RuntimeError(f"Runtime-cache target tensor {runtime_name!r} no longer exists")
+                if target.device.type != "cpu" or target.is_meta:
+                    raise RuntimeError(
+                        f"Runtime-cache target {runtime_name!r} is not an ordinary CPU tensor: {target.device}"
+                    )
+
+                if binding.file_path not in file_cache:
+                    file_cache[binding.file_path] = safe_open(
+                        binding.file_path,
+                        framework="pt",
+                        device="cpu",
+                    )
+                tensor = file_cache[binding.file_path].get_tensor(binding.storage_key)
+                if tensor.shape != target.shape or tensor.dtype != target.dtype:
+                    raise RuntimeError(
+                        f"Runtime-cache metadata changed for {runtime_name!r}: "
+                        f"cache={tuple(tensor.shape)}/{tensor.dtype}, "
+                        f"runtime={tuple(target.shape)}/{target.dtype}"
+                    )
+                replacements.append((target, tensor))
+                mapped_sources.setdefault(binding.file_path, []).append(tensor)
+        except Exception:
+            file_cache.clear()
+            raise
+
+        # Preserve Parameter objects and their loader/runtime attributes; only
+        # their storage changes. This also keeps any references held by model
+        # code valid while dropping the private allocator-backed storage.
+        for target, tensor in replacements:
+            set_tensor_storage(target, tensor)
+
+        self._mmap_transforms_by_tensor_id.clear()
+        self._mmap_file_cache = file_cache
+        self._runtime_cache_mapped_sources = mapped_sources
+        logger.info(
+            "Remapped %d final DiT tensors from runtime cache %s",
+            len(replacements),
+            plan.runtime_layout_key[:12] if plan.runtime_layout_key else "unknown",
+        )
+
+    def _try_register_runtime_cache_mmap(self) -> bool:
+        """Register complete cache mappings when the operator budget permits."""
+        pin_limit_gib = self.config.dlo_runtime_cache_pin_limit_gib
+        if pin_limit_gib <= 0:
+            return False
+        if self.device.type != "cuda":
+            logger.warning(
+                "DLO runtime-cache registration is CUDA-only; using bounded host staging on %s",
+                self.device,
+            )
+            return False
+        if not self.config.pin_cpu_memory:
+            logger.warning("DLO runtime-cache registration requires pin_cpu_memory=True; using bounded host staging")
+            return False
+        if not self._runtime_cache_mapped_sources:
+            logger.warning("DLO runtime-cache registration found no mapped sources; using bounded host staging")
+            return False
+
+        max_bytes = int(pin_limit_gib * 1024**3)
+        started = time.perf_counter()
+        try:
+            registration = CudaHostRegistration.create(
+                self._runtime_cache_mapped_sources,
+                max_bytes=max_bytes,
+            )
+        except CudaHostRegistrationCleanupError:
+            # Falling back could unmap a range that CUDA still owns. Abort the
+            # worker so process/context teardown provides the final cleanup.
+            logger.exception("DLO runtime-cache registration rollback failed")
+            raise
+        except CudaHostRegistrationError as exc:
+            logger.warning(
+                "DLO runtime-cache direct H2D unavailable (%s); using bounded host staging",
+                exc,
+            )
+            return False
+
+        self._cuda_host_registration = registration
+        logger.info(
+            "Registered %.2f GiB of shared runtime-cache mappings in %d range(s) for direct H2D in %.3f s",
+            registration.total_bytes / 1024**3,
+            registration.region_count,
+            time.perf_counter() - started,
+        )
+        return True
+
+    def _release_registered_mmap(self) -> bool:
+        registration = self._cuda_host_registration
+        if registration is None:
+            return True
+        errors = registration.close()
+        if errors:
+            logger.error(
+                "DLO runtime-cache host unregistration failed; retaining mmap handles for retry/process exit: %s",
+                errors[:3],
+            )
+            return False
+        self._cuda_host_registration = None
+        logger.info("Unregistered DLO runtime-cache host mappings")
+        return True
+
     def _init_dp_group(self) -> None:
         """Reuse the process group initialized by parallel_state.
 
@@ -1443,17 +1655,31 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._using_mmap = self.host_weight_plan is not None
         self._using_rank_local_mmap = self._using_mmap and self.dp_size <= 1
         if self._using_mmap:
-            if self.host_weight_plan.backing_kind != "checkpoint_mmap":
+            if self.host_weight_plan.backing_kind == "checkpoint_mmap":
+                self._load_weights_via_mmap(
+                    pipeline,
+                    modules,
+                    self.host_weight_plan,
+                )
+            elif self.host_weight_plan.backing_kind == "runtime_cache":
+                if self.config.dlo_use_allgather:
+                    raise ValueError("Runtime-cache backing is no-AllGather only in the first implementation")
+                self._using_rank_local_mmap = True
+                self._load_weights_from_runtime_cache(
+                    pipeline,
+                    modules,
+                    self.host_weight_plan,
+                )
+                # Assignments above dropped the model's private tensor
+                # references. Reclaim those allocations before hook setup so
+                # steady-state PSS reflects the shared mmap backing.
+                self._cleanup_after_loading()
+            else:
                 raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}")
-            self._load_weights_via_mmap(
-                pipeline,
-                modules,
-                self.host_weight_plan,
-            )
             if self._using_rank_local_mmap:
                 logger.info(
-                    "DLO rank-local mmap storage enabled: checkpoint pages are "
-                    "node-shared; each worker owns only two bounded host staging slots"
+                    "DLO rank-local mmap storage enabled: file-backed pages are node-shared; "
+                    "transfer setup will select registered direct H2D or bounded host staging"
                 )
         else:
             remaining_meta = [
@@ -1625,6 +1851,13 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if not self._all_hook_groups:
             self.enabled = bool(self._resident_blocks)
+            if self.enabled and self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+                self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+                self._runtime_cache_mapped_sources.clear()
+                assert self._resident_layer_group is not None
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+                if self._using_registered_mmap:
+                    self._resident_layer_group._cpu_staging_buffers.clear()
             if self._using_mmap and not self.enabled:
                 self._release_mmap_handles()
             return
@@ -1640,9 +1873,18 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         unified_shard_buffers = None
         if self.dp_size > 1:
             unified_shard_buffers = self._allocate_shared_shard_buffers(all_hooks)
+
+        if self._using_rank_local_mmap and self.host_weight_plan.backing_kind == "runtime_cache":
+            self._using_registered_mmap = self._try_register_runtime_cache_mmap()
+            self._runtime_cache_mapped_sources.clear()
+            for hook in all_hooks:
+                hook.registered_mmap = self._using_registered_mmap
+            if self._resident_layer_group is not None:
+                self._resident_layer_group.registered_mmap = self._using_registered_mmap
+
         unified_cpu_staging = None
         cpu_staging_events = None
-        if self._using_rank_local_mmap:
+        if self._using_rank_local_mmap and not self._using_registered_mmap:
             unified_cpu_staging = self._allocate_shared_cpu_staging_buffers(
                 all_hooks,
                 self._resident_layer_group,
@@ -1724,6 +1966,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
     def _release_mmap_handles(self) -> None:
         """Release safetensors mmap file handles."""
         self._mmap_transforms_by_tensor_id.clear()
+        self._runtime_cache_mapped_sources.clear()
         if hasattr(self, "_mmap_file_cache"):
             self._mmap_file_cache.clear()
             del self._mmap_file_cache
@@ -1748,6 +1991,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
 
         if self._using_rank_local_mmap:
             current_omni_platform.synchronize()
+        registration_released = self._release_registered_mmap()
 
         for blocks in self._blocks:
             for block in blocks:
@@ -1763,9 +2007,11 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
         self._all_hook_groups.clear()
         self._resident_blocks.clear()
         self._resident_layer_group = None
-        self._release_mmap_handles()
+        if registration_released:
+            self._release_mmap_handles()
         self._using_mmap = False
         self._using_rank_local_mmap = False
+        self._using_registered_mmap = False
         self.enabled = False
         logger.info("Distributed layer-wise offloading disabled")
 

vllm_omni/diffusion/worker/diffusion_worker.py

@@ -805,11 +805,18 @@ class DiffusionWorker:
 
     def shutdown(self) -> None:
         """Shutdown the worker and cleanup distributed environment."""
-        if self.model_runner is not None:
-            mgr = getattr(self.model_runner, "kv_transfer_manager", None)
-            if mgr is not None:
-                mgr.shutdown_prefetch()
-        destroy_distributed_env()
+        try:
+            if self.model_runner is not None:
+                mgr = getattr(self.model_runner, "kv_transfer_manager", None)
+                try:
+                    offload_backend = getattr(self.model_runner, "offload_backend", None)
+                    if offload_backend is not None:
+                        offload_backend.disable()
+                finally:
+                    if mgr is not None:
+                        mgr.shutdown_prefetch()
+        finally:
+            destroy_distributed_env()
 
 
 class CustomPipelineWorkerExtension:

vllm_omni/engine/arg_utils.py

@@ -543,6 +543,10 @@ class OrchestratorArgs:
     enable_layerwise_offload: bool = False
     enable_distributed_layerwise_offload: bool = False
     dlo_use_allgather: bool = True
+    dlo_enable_runtime_cache: bool = False
+    dlo_runtime_cache_dir: str | None = None
+    dlo_runtime_cache_lock_timeout: float = 600.0
+    dlo_runtime_cache_pin_limit_gib: float = 0.0
     dlo_resident_layers: int = 0
     boundary_ratio: float | None = None
     flow_shift: float | None = None

vllm_omni/engine/async_omni_engine.py

@@ -1000,6 +1000,10 @@ class AsyncOmniEngine:
             "enable_layerwise_offload": kwargs.get("enable_layerwise_offload", False),
             "enable_distributed_layerwise_offload": kwargs.get("enable_distributed_layerwise_offload", False),
             "dlo_use_allgather": kwargs.get("dlo_use_allgather", True),
+            "dlo_enable_runtime_cache": kwargs.get("dlo_enable_runtime_cache", False),
+            "dlo_runtime_cache_dir": kwargs.get("dlo_runtime_cache_dir", None),
+            "dlo_runtime_cache_lock_timeout": kwargs.get("dlo_runtime_cache_lock_timeout", 600.0),
+            "dlo_runtime_cache_pin_limit_gib": kwargs.get("dlo_runtime_cache_pin_limit_gib", 0.0),
             "dlo_resident_layers": kwargs.get("dlo_resident_layers", 0),
             "enforce_eager": False if kwargs.get("enforce_eager") is None else kwargs.get("enforce_eager"),
             "diffusion_compile_granularity": (

vllm_omni/entrypoints/cli/serve.py

@@ -722,6 +722,39 @@ class OmniServeCommand(CLISubcommand):
                 "independently (including existing TP shards)."
             ),
         )
+        omni_config_group.add_argument(
+            "--dlo-enable-runtime-cache",
+            action="store_true",
+            help=(
+                "After ordinary loading, publish or join a node-local mmap cache "
+                "for final DiT weights. This first version requires "
+                "--dlo-no-use-allgather."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-dir",
+            type=str,
+            default=None,
+            help=(
+                "Shared local-disk directory for DLO runtime weights. Defaults "
+                "to ~/.cache/vllm-omni/dlo-runtime-weights."
+            ),
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-lock-timeout",
+            type=float,
+            default=600.0,
+            help="Seconds to wait for another process publishing the same runtime layout.",
+        )
+        omni_config_group.add_argument(
+            "--dlo-runtime-cache-pin-limit-gib",
+            type=float,
+            default=0.0,
+            help=(
+                "Maximum GiB of read-only runtime-cache mappings each CUDA worker "
+                "may register for direct H2D. Zero keeps bounded pinned staging."
+            ),
+        )
         omni_config_group.add_argument(
             "--dlo-resident-layers",
             type=int,