@tduncan / Nvidia-Optimization-Stack-VLLM
Code Commits Issues Pulls Cloud Agent CI/CD
main
main
py 213 lines 8.25 KB
Raw
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 """This file is a pure Python wrapper for the cudart library.
4 It avoids the need to compile a separate shared library, and is
5 convenient for use when we just need to call a few functions.
6 """
7
8 import ctypes
9 from dataclasses import dataclass
10 from typing import Any
11
12 # this line makes it possible to directly load `libcudart.so` using `ctypes`
13 import torch # noqa
14
15 import vllm.envs as envs
16 from vllm.logger import init_logger
17 from vllm.platforms import current_platform
18 from vllm.utils.system_utils import find_loaded_library
19
20 logger = init_logger(__name__)
21
22 # === export types and functions from cudart to Python ===
23 # for the original cudart definition, please check
24 # https://docs.nvidia.com/cuda/cuda-runtime-api/index.html
25
26 cudaError_t = ctypes.c_int
27 cudaMemcpyKind = ctypes.c_int
28
29
30 class cudaIpcMemHandle_t(ctypes.Structure):
31 _fields_ = [("internal", ctypes.c_byte * 128)]
32
33
34 @dataclass
35 class Function:
36 name: str
37 restype: Any
38 argtypes: list[Any]
39
40
41 class CudaRTLibrary:
42 exported_functions = [
43 # ​cudaError_t cudaSetDevice ( int device )
44 Function("cudaSetDevice", cudaError_t, [ctypes.c_int]),
45 # cudaError_t cudaDeviceSynchronize ( void )
46 Function("cudaDeviceSynchronize", cudaError_t, []),
47 # ​cudaError_t cudaDeviceReset ( void )
48 Function("cudaDeviceReset", cudaError_t, []),
49 # const char* cudaGetErrorString ( cudaError_t error )
50 Function("cudaGetErrorString", ctypes.c_char_p, [cudaError_t]),
51 # ​cudaError_t cudaMalloc ( void** devPtr, size_t size )
52 Function(
53 "cudaMalloc",
54 cudaError_t,
55 [ctypes.POINTER(ctypes.c_void_p), ctypes.c_size_t],
56 ),
57 # ​cudaError_t cudaFree ( void* devPtr )
58 Function("cudaFree", cudaError_t, [ctypes.c_void_p]),
59 # ​cudaError_t cudaMemset ( void* devPtr, int value, size_t count )
60 Function(
61 "cudaMemset", cudaError_t, [ctypes.c_void_p, ctypes.c_int, ctypes.c_size_t]
62 ),
63 # ​cudaError_t cudaMemcpy ( void* dst, const void* src, size_t count, cudaMemcpyKind kind ) # noqa
64 Function(
65 "cudaMemcpy",
66 cudaError_t,
67 [ctypes.c_void_p, ctypes.c_void_p, ctypes.c_size_t, cudaMemcpyKind],
68 ),
69 # cudaError_t cudaIpcGetMemHandle ( cudaIpcMemHandle_t* handle, void* devPtr ) # noqa
70 Function(
71 "cudaIpcGetMemHandle",
72 cudaError_t,
73 [ctypes.POINTER(cudaIpcMemHandle_t), ctypes.c_void_p],
74 ),
75 # ​cudaError_t cudaIpcOpenMemHandle ( void** devPtr, cudaIpcMemHandle_t handle, unsigned int flags ) # noqa
76 Function(
77 "cudaIpcOpenMemHandle",
78 cudaError_t,
79 [ctypes.POINTER(ctypes.c_void_p), cudaIpcMemHandle_t, ctypes.c_uint],
80 ),
81 # cudaError_t cudaHostRegister ( void* ptr, size_t size, unsigned int flags )
82 Function(
83 "cudaHostRegister",
84 cudaError_t,
85 [ctypes.c_void_p, ctypes.c_size_t, ctypes.c_uint],
86 ),
87 # cudaError_t cudaHostUnregister ( void* ptr )
88 Function("cudaHostUnregister", cudaError_t, [ctypes.c_void_p]),
89 # cudaError_t cudaGetLastError ( void )
90 Function("cudaGetLastError", cudaError_t, []),
91 ]
92
93 # https://rocm.docs.amd.com/projects/HIPIFY/en/latest/tables/CUDA_Runtime_API_functions_supported_by_HIP.html # noqa
94 cuda_to_hip_mapping = {
95 "cudaSetDevice": "hipSetDevice",
96 "cudaDeviceSynchronize": "hipDeviceSynchronize",
97 "cudaDeviceReset": "hipDeviceReset",
98 "cudaGetErrorString": "hipGetErrorString",
99 "cudaMalloc": "hipMalloc",
100 "cudaFree": "hipFree",
101 "cudaMemset": "hipMemset",
102 "cudaMemcpy": "hipMemcpy",
103 "cudaIpcGetMemHandle": "hipIpcGetMemHandle",
104 "cudaIpcOpenMemHandle": "hipIpcOpenMemHandle",
105 "cudaHostRegister": "hipHostRegister",
106 "cudaHostUnregister": "hipHostUnregister",
107 "cudaGetLastError": "hipGetLastError",
108 }
109
110 # class attribute to store the mapping from the path to the library
111 # to avoid loading the same library multiple times
112 path_to_library_cache: dict[str, Any] = {}
113
114 # class attribute to store the mapping from library path
115 # to the corresponding dictionary
116 path_to_dict_mapping: dict[str, dict[str, Any]] = {}
117
118 def __init__(self, so_file: str | None = None):
119 if so_file is None:
120 so_file = (
121 find_loaded_library(
122 "libamdhip64" if current_platform.is_rocm() else "libcudart"
123 )
124 or envs.VLLM_CUDART_SO_PATH # fallback to env var
125 )
126 assert so_file is not None, (
127 "libcudart is not loaded in the current process, "
128 "try setting VLLM_CUDART_SO_PATH"
129 )
130 if so_file not in CudaRTLibrary.path_to_library_cache:
131 lib = ctypes.CDLL(so_file)
132 CudaRTLibrary.path_to_library_cache[so_file] = lib
133 self.lib = CudaRTLibrary.path_to_library_cache[so_file]
134
135 if so_file not in CudaRTLibrary.path_to_dict_mapping:
136 _funcs = {}
137 for func in CudaRTLibrary.exported_functions:
138 f = getattr(
139 self.lib,
140 CudaRTLibrary.cuda_to_hip_mapping[func.name]
141 if current_platform.is_rocm()
142 else func.name,
143 )
144 f.restype = func.restype
145 f.argtypes = func.argtypes
146 _funcs[func.name] = f
147 CudaRTLibrary.path_to_dict_mapping[so_file] = _funcs
148 self.funcs = CudaRTLibrary.path_to_dict_mapping[so_file]
149
150 def CUDART_CHECK(self, result: cudaError_t) -> None:
151 if result != 0:
152 error_str = self.cudaGetErrorString(result)
153 raise RuntimeError(f"CUDART error: {error_str}")
154
155 def cudaGetErrorString(self, error: cudaError_t) -> str:
156 return self.funcs["cudaGetErrorString"](error).decode("utf-8")
157
158 def cudaSetDevice(self, device: int) -> None:
159 self.CUDART_CHECK(self.funcs["cudaSetDevice"](device))
160
161 def cudaDeviceSynchronize(self) -> None:
162 self.CUDART_CHECK(self.funcs["cudaDeviceSynchronize"]())
163
164 def cudaDeviceReset(self) -> None:
165 self.CUDART_CHECK(self.funcs["cudaDeviceReset"]())
166
167 def cudaMalloc(self, size: int) -> ctypes.c_void_p:
168 devPtr = ctypes.c_void_p()
169 self.CUDART_CHECK(self.funcs["cudaMalloc"](ctypes.byref(devPtr), size))
170 return devPtr
171
172 def cudaFree(self, devPtr: ctypes.c_void_p) -> None:
173 self.CUDART_CHECK(self.funcs["cudaFree"](devPtr))
174
175 def cudaMemset(self, devPtr: ctypes.c_void_p, value: int, count: int) -> None:
176 self.CUDART_CHECK(self.funcs["cudaMemset"](devPtr, value, count))
177
178 def cudaMemcpy(
179 self, dst: ctypes.c_void_p, src: ctypes.c_void_p, count: int
180 ) -> None:
181 cudaMemcpyDefault = 4
182 kind = cudaMemcpyDefault
183 self.CUDART_CHECK(self.funcs["cudaMemcpy"](dst, src, count, kind))
184
185 def cudaIpcGetMemHandle(self, devPtr: ctypes.c_void_p) -> cudaIpcMemHandle_t:
186 handle = cudaIpcMemHandle_t()
187 self.CUDART_CHECK(
188 self.funcs["cudaIpcGetMemHandle"](ctypes.byref(handle), devPtr)
189 )
190 return handle
191
192 def cudaIpcOpenMemHandle(self, handle: cudaIpcMemHandle_t) -> ctypes.c_void_p:
193 cudaIpcMemLazyEnablePeerAccess = 1
194 devPtr = ctypes.c_void_p()
195 self.CUDART_CHECK(
196 self.funcs["cudaIpcOpenMemHandle"](
197 ctypes.byref(devPtr), handle, cudaIpcMemLazyEnablePeerAccess
198 )
199 )
200 return devPtr
201
202 def cudaHostRegister(self, ptr: int, size: int, flags: int = 0) -> int:
203 """Return the raw error code instead of raising, since callers can
204 fall back to pageable memory. Drain a failure with cudaGetLastError."""
205 return self.funcs["cudaHostRegister"](ctypes.c_void_p(ptr), size, flags)
206
207 def cudaHostUnregister(self, ptr: int) -> int:
208 """Return the raw error code instead of raising."""
209 return self.funcs["cudaHostUnregister"](ctypes.c_void_p(ptr))
210
211 def cudaGetLastError(self) -> int:
212 """Return and clear the error pending on this thread."""
213 return self.funcs["cudaGetLastError"]()