From 23ab14a15bef4ff356c7d5be34afe0fa01912088 Mon Sep 17 00:00:00 2001 From: Hongsheng Liu Date: Mon, 17 Aug 2026 11:50:32 +0800 Subject: [PATCH] [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 --- .../distributed_layerwise_offload.md | 115 ++++++++-- .../distributed_layerwise_offload.md | 84 +++++-- .../model_loader/test_diffusers_loader.py | 53 +++++ .../offloader/test_cuda_host_registration.py | 205 ++++++++++++++++++ .../test_distributed_layerwise_backend.py | 89 ++++++++ tests/diffusion/test_worker_wrapper_base.py | 25 +++ .../test_async_omni_diffusion_config.py | 4 + vllm_omni/config/omni_config.py | 3 + vllm_omni/config/stage_config.py | 1 + vllm_omni/diffusion/data.py | 13 +- .../model_loader/diffusers_loader.py | 18 +- vllm_omni/diffusion/offloader/base.py | 5 + .../offloader/cuda_host_registration.py | 205 ++++++++++++++++++ .../distributed_layerwise_backend.py | 202 ++++++++++++++--- .../diffusion/worker/diffusion_worker.py | 17 +- vllm_omni/engine/arg_utils.py | 1 + vllm_omni/engine/async_omni_engine.py | 1 + vllm_omni/entrypoints/cli/serve.py | 9 + 18 files changed, 976 insertions(+), 74 deletions(-) create mode 100644 tests/diffusion/offloader/test_cuda_host_registration.py create mode 100644 vllm_omni/diffusion/offloader/cuda_host_registration.py diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md index 721ead827..5f00487e1 100644 --- a/docs/design/feature/offloader/distributed_layerwise_offload.md +++ b/docs/design/feature/offloader/distributed_layerwise_offload.md @@ -46,7 +46,9 @@ produce a direct-checkpoint mmap plan for a proven-compatible runtime layout; 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. +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 @@ -138,12 +140,13 @@ 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 and feed -the existing bounded two-slot host staging path. Because file-backed pages are -not pinned, every streamed block is first packed from 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. +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 @@ -151,6 +154,43 @@ 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 @@ -203,6 +243,10 @@ 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. Supported CPU layouts may then enter the opt-in runtime cache. @@ -225,9 +269,9 @@ This mode means: | 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; the optional runtime cache shares equivalent final layouts while excluding DP rank from its identity. | -| **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. | -| **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. | +| **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 @@ -286,6 +330,8 @@ Current source-level validation includes: 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 @@ -397,13 +443,48 @@ rose from 3.75 to 32.49 seconds, while communication-kernel time rose from summed across ranks and can overlap; they identify the mechanism, not wall time. Matching private/cache outputs were byte-identical. -The result confirms that runtime-cache v1 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 +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 - Use **DP + DLO AllGather** for the supported throughput and host-memory @@ -413,8 +494,10 @@ Phase B memory validation. - Use **no-AllGather** when independent replica execution is required. TP1 direct-mmap deployments can share checkpoint pages per node. Enable the runtime cache only when independent replicas have repeated final layouts and - host capacity is more important than the added mmap-to-pinned staging cost. - Benchmark the target CPU/memory topology, especially for matching - coordinates across multiple TP engines. + 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. diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md index 6458c8b04..eff153fe1 100644 --- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md +++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md @@ -7,7 +7,9 @@ each rank streams a complete block independently. Compatible TP1 deployments can share checkpoint-backed host pages among processes on the same node; otherwise DLO streams the ordinary loader's rank-local tensors. The optional runtime-weight cache extends node-local sharing to final ordinary-loader -layouts, including matching TP coordinates across independent engines. +layouts, including matching TP coordinates across independent engines. On +CUDA, an explicit registration budget can make those shared mappings directly +H2D-copyable and avoid recurrent mmap-to-pinned host packing. See the [DLO feature design](../../../design/feature/offloader/distributed_layerwise_offload.md) for the implementation contract and compatibility matrix. @@ -48,7 +50,8 @@ vllm serve /path/to/model --omni \ --tensor-parallel-size 2 \ --dlo-no-use-allgather \ --dlo-enable-runtime-cache \ - --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights + --dlo-runtime-cache-dir /var/cache/vllm-omni/dlo-runtime-weights \ + --dlo-runtime-cache-pin-limit-gib 80 # Sequence parallel deployment vllm serve /path/to/model --omni \ @@ -65,6 +68,7 @@ omni = Omni( dlo_use_allgather=False, dlo_enable_runtime_cache=True, dlo_runtime_cache_dir="/var/cache/vllm-omni/dlo-runtime-weights", + dlo_runtime_cache_pin_limit_gib=80, ) ``` @@ -79,6 +83,7 @@ omni = Omni( | `--dlo-enable-runtime-cache` | Normalize final ordinary-loader DiT weights into a node-local mmap cache; requires no-AllGather | `false` | | `--dlo-runtime-cache-dir PATH` | Shared local-disk cache root | `~/.cache/vllm-omni/dlo-runtime-weights` | | `--dlo-runtime-cache-lock-timeout SECONDS` | Maximum wait for another process publishing the same layout | `600` | +| `--dlo-runtime-cache-pin-limit-gib GIB` | Per-worker limit for registering the complete read-only runtime-cache mapping for direct H2D; zero keeps bounded host staging | `0` | | `--dlo-resident-layers N` | Keep N leading main-DiT blocks on device; requires no-AllGather and model-declared resident paths | `0` | ## Host-weight loading @@ -95,14 +100,21 @@ DLO may still consume those TP-local tensors, but this is a compatibility path: it does not share checkpoint-backed runtime weights across DP replicas and provides no shared-mmap host-memory guarantee. -When `--dlo-enable-runtime-cache` is set with no-AllGather, that fallback gains -a second storage opportunity. Every rank still completes ordinary loading, -loader callbacks, post-load processing, validation, and calibration. It then -hashes the final CPU DiT parameters and persistent buffers. One process per -equivalent runtime layout publishes sharded safetensors plus a checksummed -manifest; peers validate and mmap the same files. DLO replaces only tensor -storage, preserves the final `Parameter` objects, and releases the private -allocator-backed copies before installing its streaming hooks. +When `--dlo-enable-runtime-cache` is set with no-AllGather, ordinary loading +gains a second storage opportunity. Every rank completes loader callbacks, +post-load processing, validation, and calibration. It then hashes the final +CPU DiT parameters and persistent buffers. One process per equivalent runtime +layout publishes sharded safetensors plus a checksummed manifest; peers +validate and mmap the same files. DLO replaces only tensor storage, preserves +the final `Parameter` objects, and releases the private allocator-backed copies +before installing its streaming hooks. + +A positive `--dlo-runtime-cache-pin-limit-gib` deliberately selects this final +runtime layout even when a direct TP1 checkpoint plan is available. Raw +checkpoint mappings may require deferred loader transforms such as grouped-QKV +reordering; registering them would still leave recurrent host work. Publishing +the fully transformed layout once makes every streamed source directly +copyable. The runtime layout identity excludes DP rank and SP rank. It includes TP world size/rank and conservative SP implementation/world-size guards. Therefore two @@ -110,13 +122,27 @@ independent `TP=2` engines normally create two entries: both TP0 workers map one entry and both TP1 workers map the other. Cache sharing does not create a process group or an inference-time collective. -At inference time, cached tensors remain pageable mmap views. DLO packs each -streamed block into one of two bounded pinned host slots before the asynchronous -H2D copy. The private ordinary-loader path instead creates model-sized pinned -storage once during initialization. The cache therefore exchanges recurrent -host-to-host copies and DDR bandwidth for lower steady-state PSS; competing TP -ranks or engines can turn staging skew into collective wait time. Treat this as -a host-capacity mode and benchmark it on the target CPU/memory topology. +With a zero registration budget, cached tensors remain pageable mmap views. +DLO packs each streamed block into one of two bounded pinned host slots before +the asynchronous H2D copy. This lowers steady-state PSS but adds recurrent +host-to-host copies; competing TP ranks or engines can turn DDR contention and +staging skew into collective wait time. + +With a positive budget on CUDA, DLO instead attempts to register the complete +page-aligned cache mapping as read-only pinned host memory. Registration is +all-or-nothing: the total mapped range must fit the per-worker budget, and any +CUDA registration or verification failure rolls back partial registrations and +uses the bounded staging path. If CUDA cannot roll back a partial registration, +worker initialization aborts instead of unmapping memory that CUDA may still +own. On success, each immutable tensor view is copied directly into the +existing rotating device buffers on the H2D stream, so no private host staging +slots are allocated. + +CUDA registration is process-local, but it does not clone file contents. The +same physical page-cache pages remain shared among equivalent workers; each +worker does, however, hold its own registration reference and must satisfy the +host's page-locking limits. Shutdown synchronizes outstanding GPU work, +unregisters every range, and only then closes the mmap handles. The mmap plan skips only dedicated DiT weight sources. Other component sources, such as a text encoder loaded through the shared diffusion loader, continue to @@ -165,6 +191,13 @@ entry is published. Operators may remove entries or the complete cache root only after all workers using its mmap files have stopped. The feature is opt-in for this reason. +For direct H2D, set the registration limit above the complete page-aligned +layout size reported in the worker log. Partial registration is intentionally +unsupported because mixing direct and staged tensors would complicate buffer +lifetime and make performance topology-dependent. Registration is CUDA-only; +an insufficient budget, unsupported runtime, OS/CUDA pinning limit, or +registration error is reported and falls back to bounded staging. + ## Declarative topology Models may declare an `OffloadPlan` instead of embedding offload logic: @@ -196,7 +229,9 @@ must enter each collective. - Direct checkpoint mmap currently requires TP1. TP greater than one falls back to the ordinary TP-aware loader. With the optional runtime cache, equivalent processes at the same TP coordinate can then share that final - rank-local layout in no-AllGather mode. + rank-local layout in no-AllGather mode. A positive registration budget also + chooses the final runtime cache at TP1 so deferred checkpoint transforms are + not repeated during inference. - HSDP plus AllGather is rejected to avoid double sharding. HSDP without AllGather has limited end-to-end validation. - Per-tensor online FP8 linears use the ordinary loader and can run with either @@ -212,16 +247,19 @@ must enter each collective. The runtime cache improves steady-state host PSS, not startup peak: every rank loads private weights before remapping, and validation adds full content-hash -passes. Skip-load-on-cache-hit and cache eviction remain follow-up work in +passes. Direct H2D removes recurrent staging but does not change that startup +contract. Skip-load-on-cache-hit and cache eviction remain follow-up work in [RFC #6195](https://github.com/vllm-project/vllm-omni/issues/6195). -It is also not latency-neutral. In the four-L20X MiniMax-H3 smoke documented in -the L20X runtime-cache matrix in the +Zero-budget staging is also not latency-neutral. In the four-L20X MiniMax-H3 +smoke documented in the L20X runtime-cache matrix in the [feature design](../../../design/feature/offloader/distributed_layerwise_offload.md#l20x-runtime-cache-smoke-matrix), two independent TP2 engines reduced request-time host PSS by 36.6% with the cache, but recurrent mmap-to-pinned staging reduced combined throughput by -69.3%. HBM and GPU compute were unchanged. Use the cache only when that -measured memory/latency tradeoff is appropriate for the deployment. +69.3%. The registered path recovered near-private per-engine latency while +retaining the shared-memory benefit; a matched TP1 run improved mean wave +latency from 33.87 seconds to 4.88 seconds with byte-identical outputs. See the +feature design for the full smoke-test caveats and metrics. See the [Cosmos3 DistOffload recipe](https://github.com/vllm-project/vllm-omni/blob/main/recipes/cosmos3/Cosmos3-DistOffload.md) for an end-to-end example. diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py index db868bf63..56557bc07 100644 --- a/tests/diffusion/model_loader/test_diffusers_loader.py +++ b/tests/diffusion/model_loader/test_diffusers_loader.py @@ -492,6 +492,59 @@ def test_dlo_runtime_cache_is_built_after_final_post_load_mutation(monkeypatch, 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 diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py new file mode 100644 index 000000000..37bdff116 --- /dev/null +++ b/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] diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py index f649565ae..e49375683 100644 --- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py +++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py @@ -546,6 +546,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): @@ -636,6 +661,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 ): @@ -995,6 +1038,52 @@ class TestMmapWeightLoading: 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): diff --git a/tests/diffusion/test_worker_wrapper_base.py b/tests/diffusion/test_worker_wrapper_base.py index 170faec80..073b379af 100644 --- a/tests/diffusion/test_worker_wrapper_base.py +++ b/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 diff --git a/tests/entrypoints/test_async_omni_diffusion_config.py b/tests/entrypoints/test_async_omni_diffusion_config.py index 1befcfe44..565eeaf84 100644 --- a/tests/entrypoints/test_async_omni_diffusion_config.py +++ b/tests/entrypoints/test_async_omni_diffusion_config.py @@ -333,6 +333,8 @@ def test_serve_cli_forwards_distributed_offload_residency(): "/var/cache/vllm-omni-dlo", "--dlo-runtime-cache-lock-timeout", "30", + "--dlo-runtime-cache-pin-limit-gib", + "80", "--dlo-resident-layers", "20", ] @@ -347,12 +349,14 @@ def test_serve_cli_forwards_distributed_offload_residency(): 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 diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py index 1ca12cd4e..f7327c5b9 100644 --- a/vllm_omni/config/omni_config.py +++ b/vllm_omni/config/omni_config.py @@ -690,6 +690,7 @@ class _DiffusionConfigProjection: 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" @@ -756,6 +757,8 @@ class _DiffusionConfigProjection: 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() diff --git a/vllm_omni/config/stage_config.py b/vllm_omni/config/stage_config.py index 5ac625bb3..9148198e6 100644 --- a/vllm_omni/config/stage_config.py +++ b/vllm_omni/config/stage_config.py @@ -385,6 +385,7 @@ class StageDeployConfig: 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 diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py index f586367d0..cb6dda5d6 100644 --- a/vllm_omni/diffusion/data.py +++ b/vllm_omni/diffusion/data.py @@ -708,13 +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 when direct - # checkpoint mmap is unavailable. The first version is no-AllGather only. + # 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 @@ -945,10 +950,14 @@ 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', " diff --git a/vllm_omni/diffusion/model_loader/diffusers_loader.py b/vllm_omni/diffusion/model_loader/diffusers_loader.py index f89b7e8ee..e0c5149fb 100644 --- a/vllm_omni/diffusion/model_loader/diffusers_loader.py +++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py @@ -441,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, @@ -456,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 diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py index ea74fbfb7..38e4c3255 100644 --- a/vllm_omni/diffusion/offloader/base.py +++ b/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, ) diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py new file mode 100644 index 000000000..e4a4ed305 --- /dev/null +++ b/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 diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py index 597f478dd..42541537d 100644 --- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py +++ b/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.""" @@ -1135,6 +1199,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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. @@ -1169,6 +1234,7 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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 @@ -1181,12 +1247,74 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): 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. @@ -1550,8 +1678,8 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend): raise ValueError(f"Unsupported DLO host-weight backing: {self.host_weight_plan.backing_kind}") if self._using_rank_local_mmap: logger.info( - "DLO rank-local mmap storage enabled: file-backed 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 = [ @@ -1723,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 @@ -1738,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, @@ -1822,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 @@ -1846,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: @@ -1861,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") diff --git a/vllm_omni/diffusion/worker/diffusion_worker.py b/vllm_omni/diffusion/worker/diffusion_worker.py index d9030e891..597165330 100644 --- a/vllm_omni/diffusion/worker/diffusion_worker.py +++ b/vllm_omni/diffusion/worker/diffusion_worker.py @@ -859,11 +859,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: diff --git a/vllm_omni/engine/arg_utils.py b/vllm_omni/engine/arg_utils.py index db79c2a34..8c4c7fc43 100644 --- a/vllm_omni/engine/arg_utils.py +++ b/vllm_omni/engine/arg_utils.py @@ -548,6 +548,7 @@ class OrchestratorArgs: 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 diff --git a/vllm_omni/engine/async_omni_engine.py b/vllm_omni/engine/async_omni_engine.py index cc1d5c469..940c45abc 100644 --- a/vllm_omni/engine/async_omni_engine.py +++ b/vllm_omni/engine/async_omni_engine.py @@ -1043,6 +1043,7 @@ class AsyncOmniEngine: "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": ( diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py index 1e74cb2a4..154374852 100644 --- a/vllm_omni/entrypoints/cli/serve.py +++ b/vllm_omni/entrypoints/cli/serve.py @@ -746,6 +746,15 @@ class OmniServeCommand(CLISubcommand): 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,