anodyine/vllm-omni
1
0 Fork 0
Code Issues 4 Pull Requests 2 Actions Packages Projects Releases Wiki Activity

[Diffusion] Register runtime cache for direct DLO H2D

Browse Source 
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>
This commit is contained in:
Hongsheng Liu committed 2026-08-19 11:21:28 +08:00
1 parent 67f15e73cd
commit 23ab14a15b
18 files changed
+976 -74
Download Patch File Download Diff File Expand all files Collapse all files

No files matched your search

docs/design/feature/offloader/distributed_layerwise_offload.md
+99 -16
View File
@@ -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.
docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md
+61 -23
View File
@@ -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.
tests/diffusion/model_loader/test_diffusers_loader.py
+53
View File
@@ -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
tests/diffusion/offloader/test_cuda_host_registration.py
+205
View File
@@ -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
+89
View File
@@ -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):
tests/diffusion/test_worker_wrapper_base.py
+25
View File
@@ -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
+4
View File
@@ -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
vllm_omni/config/omni_config.py
+3
View File
@@ -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()
vllm_omni/config/stage_config.py
+1
View File
@@ -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
vllm_omni/diffusion/data.py
+11 -2
View File
@@ -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', "
vllm_omni/diffusion/model_loader/diffusers_loader.py
+17 -1
View File
@@ -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
vllm_omni/diffusion/offloader/base.py
+5
View File
@@ -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
+205
View File
@@ -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
+175 -27
View File
@@ -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")
vllm_omni/diffusion/worker/diffusion_worker.py
+12 -5
View File
@@ -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:
vllm_omni/engine/arg_utils.py
+1
View File
@@ -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
vllm_omni/engine/async_omni_engine.py
+1
View File
@@ -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": (
vllm_omni/entrypoints/cli/serve.py
+9
View File
@@ -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,