| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | # isort: skip_file |
| 5 | # ruff: noqa: E402 |
| 6 | # mypy: disable-error-code="misc, assignment" |
| 7 | |
| 8 | from typing import Any |
| 9 | |
| 10 | import numpy as np |
| 11 | |
| 12 | # Patch torch APIs |
| 13 | import torch |
| 14 | |
| 15 | |
| 16 | def noop(*args: Any, **kwargs: Any) -> None: |
| 17 | pass |
| 18 | |
| 19 | |
| 20 | # Distinct no-op so empty_cache does not alias synchronize: Dynamo's |
| 21 | # handle_synchronize is keyed on that object and asserts on CPU-only hosts. |
| 22 | def empty_cache_noop(*args: Any, **kwargs: Any) -> None: |
| 23 | pass |
| 24 | |
| 25 | |
| 26 | def fake_pin_memory(self: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor: |
| 27 | return self |
| 28 | |
| 29 | |
| 30 | class _EventPlaceholder: |
| 31 | def __init__(self, *args, **kwargs) -> None: |
| 32 | self.record = noop |
| 33 | self.wait = noop |
| 34 | self.synchronize = noop |
| 35 | |
| 36 | |
| 37 | class _StreamPlaceholder: |
| 38 | def __init__(self, *args, **kwargs) -> None: |
| 39 | self.wait_stream = noop |
| 40 | self.wait_event = noop |
| 41 | self.record_event = noop |
| 42 | self.synchronize = noop |
| 43 | self.query = lambda: True |
| 44 | self.device = torch.device("cpu") |
| 45 | |
| 46 | def __enter__(self, *args, **kwargs): |
| 47 | return self |
| 48 | |
| 49 | def __exit__(self, exc_type, exc_val, exc_tb): |
| 50 | pass |
| 51 | |
| 52 | |
| 53 | from vllm.utils.cpu_resource_utils import get_memory_node_info |
| 54 | |
| 55 | |
| 56 | def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]: |
| 57 | meminfo = get_memory_node_info() |
| 58 | return meminfo.available_memory, meminfo.total_memory |
| 59 | |
| 60 | |
| 61 | torch.Event = _EventPlaceholder |
| 62 | torch.cuda.Event = _EventPlaceholder |
| 63 | torch.cuda.Stream = _StreamPlaceholder |
| 64 | torch.cuda.set_stream = noop |
| 65 | torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder() |
| 66 | torch.cuda.stream = lambda *args, **kwargs: _StreamPlaceholder() |
| 67 | torch.accelerator.synchronize = noop |
| 68 | torch.accelerator.empty_cache = empty_cache_noop |
| 69 | torch.Tensor.pin_memory = fake_pin_memory |
| 70 | torch.Tensor.record_stream = noop |
| 71 | torch.accelerator.get_memory_info = get_memory_info |
| 72 | |
| 73 | # Patch vLLM torch utils |
| 74 | import vllm.utils.torch_utils as torch_utils |
| 75 | |
| 76 | |
| 77 | def async_tensor_h2d( |
| 78 | data: list | np.ndarray | torch.Tensor, |
| 79 | device: str | torch.device | None = None, |
| 80 | dtype: torch.dtype | None = None, |
| 81 | out: torch.Tensor | None = None, |
| 82 | ) -> torch.Tensor: |
| 83 | if isinstance(data, np.ndarray): |
| 84 | data = torch.from_numpy(data) |
| 85 | if not isinstance(data, torch.Tensor): |
| 86 | data = torch.tensor(data, dtype=dtype, device="cpu") |
| 87 | elif out is None: |
| 88 | return data.to(dtype=dtype) |
| 89 | return data if out is None else out.copy_(data) |
| 90 | |
| 91 | |
| 92 | torch_utils.async_tensor_h2d = async_tensor_h2d |
| 93 | |
| 94 | # Patch model runner APIs |
| 95 | import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils |
| 96 | import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils |
| 97 | |
| 98 | gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer |
| 99 | |
| 100 | # Patch Triton |
| 101 | from vllm.triton_utils import HAS_TRITON, tl |
| 102 | |
| 103 | if HAS_TRITON: |
| 104 | tl.debug_barrier = noop |