vllm.device_allocator

Modules:

• alloc_conf –

  Helpers for reading and editing PyTorch's allocator configuration.

• cumem –
• sleep_mode_backend –

  Pluggable sleep-mode backends (RFC #34303).

• xpumem –

Functions:

• discardable_allocations –

  Weights allocations whose live tensors all start at these addresses.

discardable_allocations(allocator_and_pools, tensors)

Weights allocations whose live tensors all start at these addresses.

Source code in vllm/device_allocator/__init__.py

def discardable_allocations(
    allocator_and_pools: Mapping[str, Iterable[tuple[Any, Any]]],
    tensors: Collection[int],
) -> set[int]:
    """Weights allocations whose live tensors all start at these addresses."""
    if not tensors:
        return set()
    return {
        allocation["address"]
        for pool, _ in allocator_and_pools.get("weights", ())
        for allocation in pool.snapshot()
        if all(
            block["state"] == "inactive" or block["address"] in tensors
            for block in allocation["blocks"]
        )
    }