vllm_omni.diffusion.offloader.tensor_utils ¶

Shared tensor utilities for distributed layerwise offload.

These helpers are used by both DistributedLayerwiseOffloadHook and DistributedLayerwiseOffloadBackend, and can be reused by other offload backends.

TensorStorageSpec `dataclass` ¶

One tensor's target and physical-storage layout.

name `instance-attribute` ¶

name: str

storage_numel `instance-attribute` ¶

storage_numel: int

stride `instance-attribute` ¶

stride: tuple[int, ...]

target `instance-attribute` ¶

target: Tensor

value `instance-attribute` ¶

value: Tensor

clear_block_storage ¶

clear_block_storage(
    params: Mapping[str, Tensor],
    buffers: Mapping[str, Tensor],
    ready_event: Any | None,
) -> None

Wait for an outstanding transfer, then release one block's residency.

clear_tensor_storage ¶

clear_tensor_storage(tensors: Iterable[Tensor]) -> None

Atomically replace tensor storage with offload placeholders.

describe_tensor_storage ¶

describe_tensor_storage(
    named_tensors: Iterable[tuple[str, Tensor]],
    transforms: Mapping[int, Callable[[Tensor], Tensor]]
    | None = None,
) -> list[TensorStorageSpec]

Resolve local tensors and their transport-preserving layouts.

dtype_size ¶

dtype_size(dtype: dtype) -> int

Return element size in bytes for a torch.dtype.

flatten_physical_storage ¶

flatten_physical_storage(
    tensor: Tensor, storage_numel: int | None = None
) -> Tensor

Pack a possibly strided tensor in physical storage order.

group_named_tensors_by_dtype ¶

group_named_tensors_by_dtype(
    params: Mapping[str, Tensor],
    buffers: Mapping[str, Tensor],
) -> dict[dtype, list[tuple[str, Tensor]]]

Group parameters and buffers by dtype while preserving their order.

is_dtensor ¶

is_dtensor(t: Tensor) -> bool

Check if tensor is a DTensor.

is_materialized_tensor ¶

is_materialized_tensor(t: Tensor) -> bool

Check if tensor holds real data (not meta or empty placeholder).

make_offload_placeholder ¶

make_offload_placeholder(tensor: Tensor) -> Tensor

Create a zero-element placeholder to free GPU memory.

materialization_probe ¶

materialization_probe(
    params: Mapping[str, Tensor],
    buffers: Mapping[str, Tensor],
) -> Tensor | None

Pick one non-empty tensor that tracks a block's atomic storage state.

module_materialization_probe ¶

module_materialization_probe(
    module: Module,
) -> Tensor | None

Capture a probe before another hook releases the module's storage.

physical_storage_numel ¶

physical_storage_numel(tensor: Tensor) -> int

Return the storage span needed to preserve tensor's stride.

restore_tensor_storage ¶

restore_tensor_storage(
    target: Tensor, value: Tensor, *, device: device | str
) -> None

Detach target from an offloader-owned backing with a real copy.

set_tensor_storage ¶

set_tensor_storage(target: Tensor, value: Tensor) -> None

Replace target's underlying storage with value (zero-copy).

tensor_storage_metadata ¶

tensor_storage_metadata(
    spec: TensorStorageSpec,
    offset: int,
    *,
    include_device: bool = False,
) -> dict[str, Any]

Build reconstruction metadata shared by offload transports.