A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 7, 11:30 UTC
[CI] Fix mypy for vllm/device allocator (#35518)
Signed-off-by: Martin Hickey <martin.hickey@ie.ibm.com> Co-authored-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>
Martin Hickey committed
Mar 2, 2026 at 15:53 UTC
7560d674c9b35e4d1f1a91bfa7bbd18a949aafe0
2 files changed
+10
-8
tools/pre_commit/mypy.py
-1
@@ -41,7 +41,6 @@ EXCLUDE = [
41
# TODO: Remove these entries after fixing mypy errors.
42
"vllm/benchmarks",
43
"vllm/config",
44
- "vllm/device_allocator",
44
"vllm/reasoning",
45
"vllm/tool_parser",
46
]
vllm/device_allocator/cumem.py
+10
-7
@@ -11,7 +11,7 @@
11
import dataclasses
12
import gc
13
import os
14
-from collections.abc import Callable
14
+from collections.abc import Callable, Iterator
15
from contextlib import contextmanager
16
from typing import Any
17
@@ -25,6 +25,7 @@ logger = init_logger(__name__)
25
26
27
cumem_available = False
28
+libcudart: Any = None
29
try:
30
from vllm.cumem_allocator import (
31
init_module,
@@ -41,9 +42,7 @@ except ModuleNotFoundError:
42
init_module = None
43
python_create_and_map = None
44
python_unmap_and_release = None
44
- CudaRTLibrary = None
45
lib_name = None
46
- libcudart = None
46
47
# py_device, py_alignedSize, py_d_mem, py_p_memHandle
48
HandleType = tuple[int, int, int, int]
@@ -65,7 +64,8 @@ def unmap_and_release(allocation_handle: HandleType) -> None:
64
65
66
def get_pluggable_allocator(
68
- python_malloc_fn: Callable[[int], int], python_free_func: Callable[[int, int], None]
67
+ python_malloc_fn: Callable[[HandleType], None],
68
+ python_free_func: Callable[[int], HandleType],
69
) -> torch.cuda.memory.CUDAPluggableAllocator:
70
init_module(python_malloc_fn, python_free_func)
71
new_alloc = torch.cuda.memory.CUDAPluggableAllocator(
@@ -76,8 +76,11 @@ def get_pluggable_allocator(
76
77
@contextmanager
78
def use_memory_pool_with_allocator(
79
- python_malloc_fn: Callable[[int], int], python_free_func: Callable[[int, int], None]
80
-) -> None:
79
+ python_malloc_fn: Callable[[HandleType], None],
80
+ python_free_func: Callable[[int], HandleType],
81
+) -> Iterator[
82
+ tuple[torch.cuda.memory.MemPool, torch.cuda.memory.CUDAPluggableAllocator]
83
+]:
84
new_alloc = get_pluggable_allocator(python_malloc_fn, python_free_func)
85
mem_pool = torch.cuda.memory.MemPool(new_alloc._allocator)
86
with torch.cuda.memory.use_mem_pool(mem_pool):
@@ -109,7 +112,7 @@ class CuMemAllocator:
112
not work as expected.
113
"""
114
112
- instance: "CuMemAllocator" = None
115
+ instance: "CuMemAllocator | None" = None
116
default_tag: str = "default"
117
118
@staticmethod