vllm_omni.diffusion.offloader.module_residency ¶

Retain reusable allocator blocks without monopolizing device memory.

Component offload normally calls empty_cache after every stage. That makes the next stage return to the device allocator even though PyTorch's cached blocks are immediately reusable. This policy keeps the cache while both of these bounds hold:

• cached-but-unallocated memory is at most 25% of device capacity; and
• at least 5% of device capacity is physically free.

Missing memory telemetry is handled conservatively by releasing the cache. Failure paths can force release; normal executor shutdown keeps its own unconditional device-cache cleanup because this policy is not global.

device = device

max_cached_fraction = max_cached_fraction

min_free_fraction = min_free_fraction

release_if_needed(*, force: bool = False) -> bool

Stage immutable module groups without copying device weights to CPU.

nn.Module.to("cpu") performs a device-to-host copy for every parameter after each forward. In inference the weights are immutable, so retain one pinned CPU master instead. load materializes device storage from that master; offload only rebinds Parameters and buffers to the master.

A module iterable is treated as one staging group. It uses one copy stream and one reusable completion event. Tensors sharing storage keep their shapes, strides, offsets, dtypes, and aliases across every transition.

cache_retention = cache_retention

copy_stream = (
    copy_stream
    if copy_stream is not None
    else current_omni_platform.Stream()
)

device = device

loaded = False

load() -> None

offload() -> None

set_cache_retention(
    cache_retention: BoundedAllocatorCache | None,
) -> None