diff --git a/docs/design/feature/offloader/distributed_layerwise_offload.md b/docs/design/feature/offloader/distributed_layerwise_offload.md
index 94684cd81..e3f3e3dcf 100644
--- a/docs/design/feature/offloader/distributed_layerwise_offload.md
+++ b/docs/design/feature/offloader/distributed_layerwise_offload.md
@@ -46,13 +46,14 @@ produce a direct-checkpoint mmap plan for a proven-compatible runtime layout;
otherwise it uses the ordinary loader. In no-AllGather mode, the Phase B host
weight cache can explicitly normalize supported final ordinary-loader DiT
tensors into an immutable node-local mmap entry. Ordinary pinned weights remain
-the default after checkpoint mmap fallback because bounded cache staging can
-reduce throughput. Replicas can share either proven-compatible checkpoint pages
-or, when selected, equivalent final runtime layouts. An opt-in
-host-registration budget can make the complete final mapping
-directly H2D-copyable when the platform provides a backend, avoiding the
-recurrent host packing otherwise required by mmap. CUDA is the first
-implementation.
+the default after checkpoint mmap fallback when the cache is not selected; cache
+publication and validation add startup and disk lifecycle costs. Replicas can
+share either proven-compatible checkpoint pages or, when selected, equivalent
+final runtime layouts. Under the existing
+pinned-memory policy, the backend attempts to make the complete final mapping
+directly H2D-copyable when the platform provides registration, avoiding the
+recurrent host packing otherwise required by mmap. An optional positive limit
+adds a per-worker safety ceiling. CUDA is the first implementation.
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
@@ -107,7 +108,7 @@ flowchart TB
Discovery["ModuleDiscovery + optional OffloadPlan"] --> Backend
Contract -.->|"passed once by the runner"| Backend
- Backend -.->|"host weight cache + positive budget"| Registration["host_registration
platform-neutral lifecycle"]
+ Backend -.->|"host weight cache + pinned-memory policy"| Registration["host_registration
platform-neutral lifecycle"]
Registration --> CUDA["CUDA registration backend"]
Backend --> Hooks["DistributedLayerwiseOffloadHook
one circular hook per streamed block"]
Forward["pipeline.forward"] --> Hooks
@@ -119,7 +120,7 @@ flowchart TB
| Component | Owns | Must not own |
| --- | --- | --- |
| CLI, engine arguments, and `OmniDiffusionConfig` | User policy, defaults, validation, and propagation to every worker | Checkpoint inspection, cache publication, or buffer allocation |
-| `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, and registration budget | Host-weight compatibility decisions |
+| `OffloadConfig.from_od_config()` | Offload-strategy selection, effective DLO group size, AllGather policy, resident-layer count, pinned-memory policy, and optional registration ceiling | Host-weight compatibility decisions |
| `DiffusersPipelineLoader` | Ordinary loading lifecycle and the decision to return an optional `HostWeightPlan` | DLO process groups or block-stream scheduling |
| Checkpoint mmap planner | Proving that checkpoint tensors can represent the required runtime layout before ordinary DiT materialization is skipped | Host weight cache identity or transfer policy |
| `host_weight_cache` | Normalizing final ordinary-loader tensors into an immutable node-local entry and returning a validated plan | Model mutation, host registration, H2D, or collectives |
@@ -133,11 +134,14 @@ flowchart TB
The configuration fields are routed to the component that can enforce their
contract. The cache root and lock timeout are consumed by the loader.
`dlo_use_host_weight_cache` explicitly selects final-layout cache publication
-after direct-checkpoint mmap fallback. It defaults to false so the ordinary
+instead of direct-checkpoint mmap. It defaults to false so the ordinary
pinned-loader path remains available without requiring an opt-out.
-`dlo_host_weight_cache_pin_limit_gib` is also visible while loading because a
-positive value also deliberately selects the transform-complete host weight
-cache; the backend later consumes the same budget for registration.
+With `pin_cpu_memory=True`, selecting the host weight cache deliberately chooses
+the transform-complete layout even when direct checkpoint mmap is available,
+because only that final layout can be registered without deferred transforms.
+The backend consumes `dlo_host_weight_cache_pin_limit_gib` as an optional
+registration ceiling; zero applies no additional ceiling and does not select
+the cache by itself.
`dlo_use_allgather` controls loader plan routing because the host weight cache
is no-AllGather-only, and it selects the backend transfer protocol. This is
coordination through typed configuration and `HostWeightPlan`, not through
@@ -246,19 +250,22 @@ sequenceDiagram
participant P as platform registration
participant H as block hooks
- L->>L: checkpoint mmap preflight
- alt loader selects a compatible checkpoint plan
- L-->>B: checkpoint plan via runner
- else ordinary loading
+ alt host weight cache explicitly selected
L->>L: load, mutate, validate, and calibrate final tensors
- opt host weight cache enabled
- L->>C: publish or join final-layout entry
- C-->>L: validated host weight cache plan or fail-closed fallback
+ L->>C: publish or join final-layout entry
+ C-->>L: validated host weight cache plan or fail-closed fallback
+ L-->>B: cache plan or materialized fallback via runner
+ else automatic checkpoint selection
+ L->>L: checkpoint mmap preflight
+ alt loader selects a compatible checkpoint plan
+ L-->>B: checkpoint plan via runner
+ else ordinary loading
+ L->>L: load, mutate, validate, and calibrate final tensors
+ L-->>B: materialized tensors via runner
end
- L-->>B: plan or materialized tensors via runner
end
B->>B: consume complete rank-local block sources
- opt host weight cache plan and positive registration budget
+ opt host weight cache plan and pinned-memory policy
B->>P: register the complete immutable mapping
P-->>B: direct-H2D source or bounded-staging fallback
end
@@ -267,7 +274,7 @@ sequenceDiagram
alt registered host weight cache mapping
H->>H: copy complete next block directly to device
else pageable mmap mapping
- H->>H: pack into a pinned staging slot, then copy to device
+ H->>H: pack into a bounded staging slot, then copy to device
else ordinary runtime tensors
H->>H: copy complete next block to device
end
@@ -287,8 +294,8 @@ the source of the complete-block H2D copy.
| Ordinary runtime tensors | AllGather group larger than one | Each rank retains a pinned local weight shard | shard H2D → DLO AllGather → device slot |
| Checkpoint mmap | AllGather group larger than one | Each rank prepares its persistent shard, then releases the source mapping | shard H2D → DLO AllGather → device slot |
| Ordinary runtime tensors | Rank-local / no-AllGather | Each rank retains the complete loader-produced local layout | complete-block H2D → device slot |
-| Checkpoint or host weight cache mmap, zero registration budget | Rank-local / no-AllGather | Immutable mapping plus two bounded pinned host-staging slots | mmap views → pinned staging slot → device slot |
-| Host weight cache mmap, positive registration budget | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot |
+| Checkpoint mmap, or host weight cache with pinned memory disabled/registration unavailable | Rank-local / no-AllGather | Immutable mapping plus two bounded host-staging slots | mmap views → staging slot → device slot |
+| Host weight cache mmap with pinned memory enabled and registration available | Rank-local / no-AllGather | Complete mapping is registered by the platform backend | registered mmap views → device slot |
The three axes therefore remain independent:
@@ -453,14 +460,15 @@ 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 host weight 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.
+For a host weight cache mapping, the default pinned-memory policy attempts to
+replace those host slots with direct copies from the shared final tensor views.
+Disabling pinned memory or a failed registration retains 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. Those ordinary pinned tensors remain the default. When the host
-weight cache or a positive registration budget is explicitly selected,
+When direct mmap preflight fails or the host weight cache is selected, the
+regular model loader remains responsible for preparing each rank's weights,
+including TP-local tensors or HSDP-managed parameters. Ordinary pinned tensors
+remain the default without cache selection. When the host weight cache is
+explicitly selected,
supported CPU layouts enter the cache; unsupported layouts keep one private
runtime copy per process.
@@ -512,13 +520,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. 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.
+streaming hooks. The source mappings remain open through DLO disable. With
+pinned memory disabled, or when registration is unavailable, they feed the
+existing bounded two-slot host-staging path. Every streamed block is first
+packed from pageable mmap storage into a staging slot 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 steady-state-oriented. It requires roughly one
local-disk model copy per runtime identity, has no eviction, and still performs
@@ -529,19 +537,20 @@ modeled safely.
### Registered host weight cache transfer
-`dlo_host_weight_cache_pin_limit_gib > 0` opts a worker into registration of
-the complete host weight cache mapping when its platform supports it. It also
-changes loader selection: no-AllGather uses the ordinary loader and
-final-layout host weight cache even when a TP1 direct-checkpoint plan is
-available. A checkpoint binding can carry a deferred adapter such as
+Selecting the host weight cache with `pin_cpu_memory=True` opts a worker into
+registration of the complete mapping when its platform supports it. It also
+changes loader selection: no-AllGather uses the ordinary loader and final-layout
+host weight 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 host weight 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 platform mutation. The CUDA backend
+adjacent ranges within that mapping. When a positive per-worker ceiling is
+configured, the complete byte count is checked against it before any platform
+mutation. The CUDA backend
first requires `cudaDevAttrHostRegisterReadOnlySupported`, then registers every
range with the read-only host-registration flag and verifies that PyTorch
recognizes every mapped tensor as pinned. A device without that capability
@@ -642,10 +651,10 @@ artifacts before making hardware or performance claims. Track that work in
direct-mmap deployments can share checkpoint pages per node. Enable the
host weight cache only when independent replicas have repeated final layouts
and benchmark the target CPU/memory topology, especially for matching
- coordinates across multiple TP engines. On CUDA, provide an explicit
- registration budget large enough for the complete final layout when
- recurrent staging cost is unacceptable; leave it at zero when page-locking
- that layout is not operationally acceptable.
+ coordinates across multiple TP engines. With the default pinned-memory
+ policy, supported CUDA workers prefer registered direct H2D. Configure a
+ positive ceiling when an explicit page-locking bound is required, or disable
+ pinned memory when pageable staging is operationally preferable.
### HSDP
diff --git a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md
index 1c27509ad..b7a7b971c 100644
--- a/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md
+++ b/docs/user_guide/diffusion/offloader/distributed_layerwise_offload.md
@@ -38,17 +38,18 @@ vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--dlo-no-use-allgather
-# Share final host weights with bounded pinned staging
+# Share final host weights; prefer registered direct H2D under the pinned-memory policy
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--dlo-no-use-allgather \
--dlo-use-host-weight-cache
-# Share final host weights and register up to 80 GiB per worker for direct H2D
+# Apply an 80 GiB per-worker safety ceiling to the same registration attempt
vllm serve /path/to/model --omni \
--enable-distributed-layerwise-offload \
--tensor-parallel-size 2 \
--dlo-no-use-allgather \
+ --dlo-use-host-weight-cache \
--dlo-host-weight-cache-pin-limit-gib 80
```
@@ -58,15 +59,18 @@ vllm serve /path/to/model --omni \
| --- | --- | --- |
| `--enable-distributed-layerwise-offload` | Enable DLO | `false` |
| `--dlo-no-use-allgather` | Stream complete rank-local blocks independently | `false` |
-| `--dlo-use-host-weight-cache` | Use the final-layout cache after checkpoint mmap fallback | `false` |
-| `--dlo-host-weight-cache-pin-limit-gib GIB` | Per-worker registration budget; positive values also select the cache | `0` |
+| `--dlo-use-host-weight-cache` | Use the final-layout cache and prefer registration when pinned memory is enabled | `false` |
+| `--dlo-host-weight-cache-pin-limit-gib GIB` | Optional per-worker registration ceiling; zero means no additional ceiling | `0` |
| `--dlo-resident-layers N` | Keep eligible leading DiT blocks resident in HBM | `0` |
No-AllGather DLO automatically uses a compatible checkpoint mmap plan. When
checkpoint mmap is unavailable, it preserves ordinary pinned loader weights by
-default. `--dlo-use-host-weight-cache` explicitly selects the final-layout
-cache with bounded staging; a positive registration budget also selects the
-cache and attempts direct H2D. The cache uses
+default. `--dlo-use-host-weight-cache` explicitly selects the transform-complete
+final-layout cache. Under the default `pin_cpu_memory=True` policy, DLO attempts
+to register the complete immutable mapping for direct H2D. Setting
+`pin_cpu_memory=False` through structured or programmatic configuration retains
+the staged mmap path instead. A positive pin limit is only an optional safety
+ceiling and does not select the cache. The cache uses
`~/.cache/vllm-omni/dlo-host-weights`. Programmatic configuration may override
`dlo_host_weight_cache_dir` and the writer lock timeout; these advanced storage
controls are intentionally not separate CLI flags.
@@ -77,15 +81,16 @@ controls are intentionally not separate CLI flags.
directory. Do not use tmpfs or a cross-node filesystem for this version.
- Cache entries are immutable and validated before use. A cache or registration
failure keeps the ordinary loader weights or falls back to bounded staging.
-- The bounded-staging cache is a memory-first mode: recurrent pageable-mmap to
+- The bounded-staging fallback is a memory-first mode: recurrent pageable-mmap to
pinned-slot copies can reduce throughput, especially across concurrent TP
- engines. Keep ordinary pinned weights unless the host-memory saving is worth
- that tradeoff, or use a validated registered path.
-- A positive registration budget must cover the complete page-aligned mapping
- reported in the worker log. Registration is all-or-nothing.
+ engines. Use it intentionally by disabling pinned host memory, or when the
+ active platform cannot register the immutable mapping.
+- When a positive registration ceiling is configured, it must cover the
+ complete page-aligned mapping reported in the worker log. Registration is
+ all-or-nothing. Zero applies no additional ceiling.
- CUDA is the first registration backend and requires read-only host
registration support for immutable cache mappings. Platforms or devices
- without that capability continue to use bounded staging; a positive budget
+ without that capability continue to use bounded staging; pinned-memory policy
requests direct H2D but does not guarantee it.
- Registration is process-local but does not duplicate the underlying file
pages. Each worker must still satisfy its platform and OS page-locking limits.
diff --git a/tests/config/test_omni_config.py b/tests/config/test_omni_config.py
index 083ac90b0..ab9935791 100644
--- a/tests/config/test_omni_config.py
+++ b/tests/config/test_omni_config.py
@@ -933,6 +933,15 @@ def test_structured_diffusion_config_restricts_host_weight_cache_to_no_allgather
)
+def test_structured_diffusion_config_requires_cache_for_registration_ceiling():
+ with pytest.raises(ValueError, match="requires dlo_use_host_weight_cache=True"):
+ omni_config_module._DiffusionConfigProjection(
+ enable_distributed_layerwise_offload=True,
+ dlo_use_allgather=False,
+ dlo_host_weight_cache_pin_limit_gib=1.0,
+ )
+
+
@pytest.mark.parametrize(
"field_name",
[
diff --git a/tests/diffusion/model_loader/test_diffusers_loader.py b/tests/diffusion/model_loader/test_diffusers_loader.py
index 029c95c61..b9f5f85f9 100644
--- a/tests/diffusion/model_loader/test_diffusers_loader.py
+++ b/tests/diffusion/model_loader/test_diffusers_loader.py
@@ -498,7 +498,8 @@ def test_dlo_host_weight_cache_is_built_after_final_post_load_mutation(monkeypat
assert loader.take_host_weight_plan() is runtime_plan
-def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(monkeypatch, tmp_path):
+@pytest.mark.parametrize("pin_cpu_memory", [True, False])
+def test_dlo_host_weight_cache_prefers_final_layout_over_checkpoint_plan(monkeypatch, tmp_path, pin_cpu_memory):
import vllm_omni.diffusion.model_loader.diffusers_loader as loader_mod
import vllm_omni.diffusion.model_loader.host_weight_cache as cache_mod
@@ -515,8 +516,10 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m
quantization_config=None,
enable_distributed_layerwise_offload=True,
dlo_use_allgather=False,
+ dlo_use_host_weight_cache=True,
+ pin_cpu_memory=pin_cpu_memory,
dlo_host_weight_cache_dir=str(tmp_path),
- dlo_host_weight_cache_pin_limit_gib=1.0,
+ dlo_host_weight_cache_pin_limit_gib=0.0,
model="unused",
)
loader = DiffusersPipelineLoader(LoadConfig(), od_config)
@@ -536,7 +539,7 @@ def test_dlo_registration_prefers_final_host_weight_cache_over_checkpoint_plan(m
loader._apply_skip_softmax_calibration = lambda _model: None # type: ignore[method-assign]
def unexpected_checkpoint_plan(*_args, **_kwargs):
- pytest.fail("registered host weight cache mode must not select the raw checkpoint mmap plan")
+ pytest.fail("selected host weight cache mode must not select the raw checkpoint mmap plan")
def build_cache(_pipeline, **_kwargs):
events.append("cache")
diff --git a/tests/diffusion/offloader/test_cuda_host_registration.py b/tests/diffusion/offloader/test_cuda_host_registration.py
index 3510de5cc..745628782 100644
--- a/tests/diffusion/offloader/test_cuda_host_registration.py
+++ b/tests/diffusion/offloader/test_cuda_host_registration.py
@@ -119,6 +119,20 @@ def test_platform_factory_dispatches_cuda(monkeypatch: pytest.MonkeyPatch) -> No
assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=4096) is sentinel
+def test_platform_factory_allows_registration_without_additional_budget(monkeypatch: pytest.MonkeyPatch) -> None:
+ sentinel = object()
+ sources = {"weights": []}
+
+ def create(actual_sources, *, max_bytes):
+ assert actual_sources is sources
+ assert max_bytes is None
+ return sentinel
+
+ monkeypatch.setattr(CudaHostRegistration, "create", staticmethod(create))
+
+ assert register_host_mappings(sources, device=torch.device("cuda"), max_bytes=None) is sentinel
+
+
def test_coalesce_ranges_aligns_and_merges_overlapping_pages() -> None:
assert _coalesce_ranges(
[(0x1003, 4096), (0x2800, 1024), (0x9001, 1)],
@@ -143,6 +157,21 @@ def test_registration_rejects_over_budget_before_calling_cuda(monkeypatch: pytes
assert runtime.registered == []
+def test_registration_without_additional_budget_registers_complete_mapping(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)
+
+ registration = CudaHostRegistration.create(
+ {"weights": [_FakeTensor(0x1003, 4096)]}, # type: ignore[list-item]
+ max_bytes=None,
+ )
+
+ assert registration.total_bytes == 8192
+ assert len(runtime.registered) == 1
+ assert registration.close() == []
+
+
def test_registration_rejects_unsupported_read_only_mapping(monkeypatch: pytest.MonkeyPatch) -> None:
monkeypatch.setattr(registration_module.torch.cuda, "is_available", lambda: True)
runtime = _FakeRuntime([0], read_only_registration_supported=False)
diff --git a/tests/diffusion/offloader/test_distributed_layerwise_backend.py b/tests/diffusion/offloader/test_distributed_layerwise_backend.py
index 3a735a636..c03e58c1d 100644
--- a/tests/diffusion/offloader/test_distributed_layerwise_backend.py
+++ b/tests/diffusion/offloader/test_distributed_layerwise_backend.py
@@ -1207,6 +1207,68 @@ class TestMmapWeightLoading:
assert events == ["unregister", "mmap"]
+ def test_host_weight_cache_registration_defaults_to_pinned_memory_policy(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ patched_offload_runtime,
+ ):
+ class Registration:
+ total_bytes = 4096
+ region_count = 1
+
+ @staticmethod
+ def close() -> list[str]:
+ return []
+
+ registration = Registration()
+ create_calls: list[tuple[dict[str, list[torch.Tensor]], int | None]] = []
+
+ def register(sources, *, device, max_bytes):
+ assert device == torch.device("cuda")
+ create_calls.append((sources, max_bytes))
+ return registration
+
+ monkeypatch.setattr(dist_backend_module, "register_host_mappings", register)
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=True,
+ dlo_use_allgather=False,
+ dlo_host_weight_cache_pin_limit_gib=0.0,
+ ),
+ torch.device("cuda"),
+ )
+ sources = {"runtime.safetensors": [torch.ones(1)]}
+ backend._host_weight_cache_mapped_sources = sources
+
+ assert backend._try_register_host_weight_cache_mmap()
+ assert create_calls == [(sources, None)]
+ assert backend._host_registration is registration
+
+ def test_host_weight_cache_registration_respects_disabled_pinned_memory(
+ self,
+ monkeypatch: pytest.MonkeyPatch,
+ patched_offload_runtime,
+ ):
+ monkeypatch.setattr(
+ dist_backend_module,
+ "register_host_mappings",
+ lambda *_args, **_kwargs: pytest.fail("registration must follow pin_cpu_memory=False"),
+ )
+ backend = DistributedLayerwiseOffloadBackend(
+ OffloadConfig(
+ strategy=OffloadStrategy.DISTRIBUTED_LAYER_WISE,
+ pin_cpu_memory=False,
+ dlo_use_allgather=False,
+ dlo_host_weight_cache_pin_limit_gib=0.0,
+ ),
+ torch.device("cuda"),
+ )
+ backend._host_weight_cache_mapped_sources = {"runtime.safetensors": [torch.ones(1)]}
+
+ assert not backend._try_register_host_weight_cache_mmap()
+ assert backend._host_registration is None
+
def test_host_weight_cache_registration_falls_back_on_unsupported_platform(
self,
patched_offload_runtime,
diff --git a/vllm_omni/config/omni_config.py b/vllm_omni/config/omni_config.py
index 48ca8ace5..3608b1f46 100644
--- a/vllm_omni/config/omni_config.py
+++ b/vllm_omni/config/omni_config.py
@@ -765,6 +765,8 @@ class _DiffusionConfigProjection:
raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
+ if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
+ raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_cache=True")
if self.tf_model_config is None:
self.tf_model_config = TransformerConfig()
diff --git a/vllm_omni/diffusion/data.py b/vllm_omni/diffusion/data.py
index 28558f026..400de9296 100644
--- a/vllm_omni/diffusion/data.py
+++ b/vllm_omni/diffusion/data.py
@@ -708,16 +708,15 @@ 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
- # Opt into publishing/joining the final-layout host weight cache when a
- # compatible direct-checkpoint mmap plan is unavailable. A positive pin
- # limit also selects the cache because registration requires its layout.
+ # Opt into publishing/joining the transform-complete final-layout host
+ # weight cache instead of direct-checkpoint mmap.
dlo_use_host_weight_cache: bool = False
# Shared local-disk root. None selects ~/.cache/vllm-omni/dlo-host-weights.
dlo_host_weight_cache_dir: str | None = None
# Maximum wait for the per-layout POSIX writer lock.
dlo_host_weight_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.
+ # Optional maximum file-backed memory a worker may register with CUDA.
+ # Zero applies no additional ceiling; pin_cpu_memory controls registration.
dlo_host_weight_cache_pin_limit_gib: float = 0.0
# Leading main-DiT blocks kept resident by distributed layerwise offload.
dlo_resident_layers: int = 0
@@ -957,6 +956,8 @@ class OmniDiffusionConfig:
raise ValueError("dlo_host_weight_cache_pin_limit_gib requires no-AllGather distributed layerwise offload")
if self.dlo_use_host_weight_cache and (not self.enable_distributed_layerwise_offload or self.dlo_use_allgather):
raise ValueError("dlo_use_host_weight_cache requires no-AllGather distributed layerwise offload")
+ if self.dlo_host_weight_cache_pin_limit_gib and not self.dlo_use_host_weight_cache:
+ raise ValueError("dlo_host_weight_cache_pin_limit_gib requires dlo_use_host_weight_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 874da9b2c..02cedefa4 100644
--- a/vllm_omni/diffusion/model_loader/diffusers_loader.py
+++ b/vllm_omni/diffusion/model_loader/diffusers_loader.py
@@ -441,15 +441,13 @@ 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_host_weight_cache = (
- _dist_offload
- and not _use_ag
- and float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0
+ _select_host_weight_cache = (
+ _dist_offload and not _use_ag and bool(getattr(self.od_config, "dlo_use_host_weight_cache", False))
)
plan_result = None
weight_sources = self._get_weight_sources(model)
- if _dist_offload and not _prefer_registered_host_weight_cache:
+ if _dist_offload and not _select_host_weight_cache:
modules = ModuleDiscovery.discover(model)
plan_result = build_checkpoint_mmap_plan(
model,
@@ -461,14 +459,14 @@ class DiffusersPipelineLoader:
online_quantization=_has_online_quant,
)
self.host_weight_plan = plan_result.plan
- elif _prefer_registered_host_weight_cache:
+ elif _select_host_weight_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 host weight cache registration requested; using the ordinary loader "
+ "DLO host weight cache selected; using the ordinary loader "
"once to publish a transform-complete final mmap layout"
)
@@ -550,10 +548,7 @@ class DiffusersPipelineLoader:
return
if getattr(self.od_config, "dlo_use_allgather", True):
return
- if not (
- getattr(self.od_config, "dlo_use_host_weight_cache", False)
- or float(getattr(self.od_config, "dlo_host_weight_cache_pin_limit_gib", 0.0)) > 0
- ):
+ if not getattr(self.od_config, "dlo_use_host_weight_cache", False):
logger.info(
"DLO host weight cache not selected; retaining ordinary host tensors after checkpoint mmap fallback"
)
diff --git a/vllm_omni/diffusion/offloader/base.py b/vllm_omni/diffusion/offloader/base.py
index 433909afc..7b02c75d2 100644
--- a/vllm_omni/diffusion/offloader/base.py
+++ b/vllm_omni/diffusion/offloader/base.py
@@ -31,8 +31,8 @@ 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 host weight cache mappings with CUDA.
- # Zero retains bounded host staging.
+ # Optional per-worker ceiling for registering shared host weight cache mappings.
+ # Zero applies no additional ceiling; pin_cpu_memory controls registration.
dlo_host_weight_cache_pin_limit_gib: float = 0.0
@classmethod
diff --git a/vllm_omni/diffusion/offloader/cuda_host_registration.py b/vllm_omni/diffusion/offloader/cuda_host_registration.py
index 1183c9a80..a52596d72 100644
--- a/vllm_omni/diffusion/offloader/cuda_host_registration.py
+++ b/vllm_omni/diffusion/offloader/cuda_host_registration.py
@@ -182,13 +182,11 @@ class CudaHostRegistration:
cls,
sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
*,
- max_bytes: int,
+ max_bytes: int | None,
) -> CudaHostRegistration:
- if max_bytes <= 0:
- raise HostRegistrationBudgetError("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:
+ if max_bytes is not None and total_bytes > max_bytes:
raise HostRegistrationBudgetError(
f"mapped host ranges need {total_bytes} bytes, exceeding the {max_bytes}-byte registration budget"
)
diff --git a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
index de298794b..3105496d9 100644
--- a/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
+++ b/vllm_omni/diffusion/offloader/distributed_layerwise_backend.py
@@ -1269,20 +1269,16 @@ class DistributedLayerwiseOffloadBackend(OffloadBackend):
)
def _try_register_host_weight_cache_mmap(self) -> bool:
- """Register complete cache mappings when the operator budget permits."""
+ """Register complete cache mappings under the pinned-memory policy."""
pin_limit_gib = self.config.dlo_host_weight_cache_pin_limit_gib
- if pin_limit_gib <= 0:
- return False
if not self.config.pin_cpu_memory:
- logger.warning(
- "DLO host weight cache registration requires pin_cpu_memory=True; using bounded host staging"
- )
+ logger.info("DLO host weight cache registration disabled by pin_cpu_memory=False; using host staging")
return False
if not self._host_weight_cache_mapped_sources:
logger.warning("DLO host weight cache registration found no mapped sources; using bounded host staging")
return False
- max_bytes = int(pin_limit_gib * 1024**3)
+ max_bytes = int(pin_limit_gib * 1024**3) if pin_limit_gib > 0 else None
started = time.perf_counter()
try:
registration = register_host_mappings(
diff --git a/vllm_omni/diffusion/offloader/host_registration.py b/vllm_omni/diffusion/offloader/host_registration.py
index 72f44f218..51c0d05d9 100644
--- a/vllm_omni/diffusion/offloader/host_registration.py
+++ b/vllm_omni/diffusion/offloader/host_registration.py
@@ -38,12 +38,13 @@ def register_host_mappings(
sources_by_mapping: Mapping[str, Sequence[torch.Tensor]],
*,
device: torch.device,
- max_bytes: int,
+ max_bytes: int | None,
) -> HostRegistration:
"""Register mappings with the active platform or report unsupported use.
CUDA is the first implementation. Other platforms retain DLO's bounded
host-staging path until they provide an equivalent registration backend.
+ ``max_bytes=None`` applies no additional ceiling to the complete mapping.
"""
if device.type == "cuda":
from .cuda_host_registration import CudaHostRegistration
diff --git a/vllm_omni/entrypoints/cli/serve.py b/vllm_omni/entrypoints/cli/serve.py
index 41ba2fa82..839d85da9 100644
--- a/vllm_omni/entrypoints/cli/serve.py
+++ b/vllm_omni/entrypoints/cli/serve.py
@@ -726,8 +726,9 @@ class OmniServeCommand(CLISubcommand):
"--dlo-use-host-weight-cache",
action="store_true",
help=(
- "Use the node-local final-layout host weight cache when direct checkpoint mmap is unavailable. "
- "This trades recurrent host staging work for shared host pages unless direct registration succeeds."
+ "Use the node-local final-layout host weight cache for no-AllGather DLO. "
+ "With the default pinned-memory policy, supported platforms register the immutable mapping "
+ "for direct H2D; otherwise DLO uses bounded host staging."
),
)
omni_config_group.add_argument(
@@ -735,9 +736,9 @@ class OmniServeCommand(CLISubcommand):
type=float,
default=0.0,
help=(
- "Maximum GiB of read-only host weight cache mappings each worker may "
- "register for direct H2D when the platform supports it. Zero keeps "
- "bounded pinned staging."
+ "Optional maximum GiB of read-only host weight cache mappings each worker may "
+ "register for direct H2D. Zero applies no additional registration ceiling; "
+ "this option requires --dlo-use-host-weight-cache."
),
)
omni_config_group.add_argument(