105 lines
2.8 KiB
Python
105 lines
2.8 KiB
Python
# SPDX-License-Identifier: Apache-2.0
|
|
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
|
|
|
|
# isort: skip_file
|
|
# ruff: noqa: E402
|
|
# mypy: disable-error-code="misc, assignment"
|
|
|
|
from typing import Any
|
|
|
|
import numpy as np
|
|
|
|
# Patch torch APIs
|
|
import torch
|
|
|
|
|
|
def noop(*args: Any, **kwargs: Any) -> None:
|
|
pass
|
|
|
|
|
|
# Distinct no-op so empty_cache does not alias synchronize: Dynamo's
|
|
# handle_synchronize is keyed on that object and asserts on CPU-only hosts.
|
|
def empty_cache_noop(*args: Any, **kwargs: Any) -> None:
|
|
pass
|
|
|
|
|
|
def fake_pin_memory(self: torch.Tensor, *args: Any, **kwargs: Any) -> torch.Tensor:
|
|
return self
|
|
|
|
|
|
class _EventPlaceholder:
|
|
def __init__(self, *args, **kwargs) -> None:
|
|
self.record = noop
|
|
self.wait = noop
|
|
self.synchronize = noop
|
|
|
|
|
|
class _StreamPlaceholder:
|
|
def __init__(self, *args, **kwargs) -> None:
|
|
self.wait_stream = noop
|
|
self.wait_event = noop
|
|
self.record_event = noop
|
|
self.synchronize = noop
|
|
self.query = lambda: True
|
|
self.device = torch.device("cpu")
|
|
|
|
def __enter__(self, *args, **kwargs):
|
|
return self
|
|
|
|
def __exit__(self, exc_type, exc_val, exc_tb):
|
|
pass
|
|
|
|
|
|
from vllm.utils.cpu_resource_utils import get_memory_node_info
|
|
|
|
|
|
def get_memory_info(*args: Any, **kwargs: Any) -> tuple[int, int]:
|
|
meminfo = get_memory_node_info()
|
|
return meminfo.available_memory, meminfo.total_memory
|
|
|
|
|
|
torch.Event = _EventPlaceholder
|
|
torch.cuda.Event = _EventPlaceholder
|
|
torch.cuda.Stream = _StreamPlaceholder
|
|
torch.cuda.set_stream = noop
|
|
torch.cuda.current_stream = lambda *args, **kwargs: _StreamPlaceholder()
|
|
torch.cuda.stream = lambda *args, **kwargs: _StreamPlaceholder()
|
|
torch.accelerator.synchronize = noop
|
|
torch.accelerator.empty_cache = empty_cache_noop
|
|
torch.Tensor.pin_memory = fake_pin_memory
|
|
torch.Tensor.record_stream = noop
|
|
torch.accelerator.get_memory_info = get_memory_info
|
|
|
|
# Patch vLLM torch utils
|
|
import vllm.utils.torch_utils as torch_utils
|
|
|
|
|
|
def async_tensor_h2d(
|
|
data: list | np.ndarray | torch.Tensor,
|
|
device: str | torch.device | None = None,
|
|
dtype: torch.dtype | None = None,
|
|
out: torch.Tensor | None = None,
|
|
) -> torch.Tensor:
|
|
if isinstance(data, np.ndarray):
|
|
data = torch.from_numpy(data)
|
|
if not isinstance(data, torch.Tensor):
|
|
data = torch.tensor(data, dtype=dtype, device="cpu")
|
|
elif out is None:
|
|
return data.to(dtype=dtype)
|
|
return data if out is None else out.copy_(data)
|
|
|
|
|
|
torch_utils.async_tensor_h2d = async_tensor_h2d
|
|
|
|
# Patch model runner APIs
|
|
import vllm.v1.worker.gpu.buffer_utils as gpu_buffer_utils
|
|
import vllm.v1.worker.cpu.buffer_utils as cpu_buffer_utils
|
|
|
|
gpu_buffer_utils.UvaBuffer = cpu_buffer_utils.UvaBuffer
|
|
|
|
# Patch Triton
|
|
from vllm.triton_utils import HAS_TRITON, tl
|
|
|
|
if HAS_TRITON:
|
|
tl.debug_barrier = noop
|