| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import dataclasses |
| 5 | from contextlib import AbstractContextManager |
| 6 | from typing import Protocol, TypeAlias |
| 7 | |
| 8 | import torch |
| 9 | |
| 10 | from vllm.platforms import current_platform |
| 11 | |
| 12 | # py_device, py_size_or_aligned_size, py_ptr, py_handle |
| 13 | # py_handle has type list[int] on ROCm and int otherwise |
| 14 | HandleType: TypeAlias = tuple[int, int, int, list[int] | int] |
| 15 | |
| 16 | # Selective wake may defer only these; any wake restores every other tag. |
| 17 | DEFERRABLE_TAGS: tuple[str, ...] = ("weights", "kv_cache") |
| 18 | |
| 19 | |
| 20 | @dataclasses.dataclass |
| 21 | class AllocationData: |
| 22 | handle: HandleType |
| 23 | tag: str |
| 24 | cpu_backup_tensor: torch.Tensor | None = None |
| 25 | is_asleep: bool = False |
| 26 | |
| 27 | |
| 28 | class MemAllocator(Protocol): |
| 29 | def use_memory_pool(self, tag: str | None = None) -> AbstractContextManager: ... |
| 30 | |
| 31 | def sleep(self, offload_tags: tuple[str, ...] | str | None = None) -> None: ... |
| 32 | |
| 33 | def discard(self, tags: tuple[str, ...] | str) -> None: ... |
| 34 | |
| 35 | def wake_up(self, tags: list[str] | None = None) -> None: ... |
| 36 | |
| 37 | def get_current_usage(self) -> int: ... |
| 38 | |
| 39 | |
| 40 | def get_mem_allocator_instance() -> MemAllocator: |
| 41 | if current_platform.is_cuda_alike(): |
| 42 | from vllm.device_allocator.cumem import CuMemAllocator |
| 43 | |
| 44 | return CuMemAllocator.get_instance() |
| 45 | |
| 46 | if current_platform.is_xpu(): |
| 47 | from vllm.device_allocator.xpumem import XpuMemAllocator |
| 48 | |
| 49 | return XpuMemAllocator.get_instance() |
| 50 | |
| 51 | raise RuntimeError( |
| 52 | "Sleep mode allocator is not available on platform " |
| 53 | f"{type(current_platform).__name__} " |
| 54 | f"(device_type={current_platform.device_type})." |
| 55 | ) |