@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / dc372b9c8a

Fork 0

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 15, 09:45 UTC

Update deprecated type hinting in `vllm/device_allocator` and `vllm/distributed` (#18126)

Signed-off-by: Harry Mellor <19981378+hmellor@users.noreply.github.com>

Harry Mellor committed May 14, 2025 at 12:07 UTC dc372b9c8aa97b5da5d4049cdccdaccef950d499

23 files changed +105 -105

pyproject.toml

-2

@@ -74,8 +74,6 @@ exclude = [

74 # Python 3.8 typing. TODO: Remove these excludes after v1.0.0

75 "vllm/attention/**/*.py" = ["UP006", "UP035"]

76 "vllm/core/**/*.py" = ["UP006", "UP035"]

77 -"vllm/device_allocator/**/*.py" = ["UP006", "UP035"]

78 -"vllm/distributed/**/*.py" = ["UP006", "UP035"]

77 "vllm/engine/**/*.py" = ["UP006", "UP035"]

78 "vllm/executor/**/*.py" = ["UP006", "UP035"]

79 "vllm/model_executor/model_loader/**/*.py" = ["UP006", "UP035"]

vllm/device_allocator/cumem.py

+5 -5

@@ -11,7 +11,7 @@ import dataclasses

11 import gc

12 import os

13 from contextlib import contextmanager

14 -from typing import Any, Callable, Dict, Optional, Tuple, Union

14 +from typing import Any, Callable, Optional, Union

16 import torch

@@ -63,7 +63,7 @@ except ModuleNotFoundError:

63 libcudart = None

65 # py_device, py_alignedSize, py_d_mem, py_p_memHandle

66 -HandleType = Tuple[int, int, int, int]

66 +HandleType = tuple[int, int, int, int]

69 @dataclasses.dataclass

@@ -148,9 +148,9 @@ class CuMemAllocator:

148 "Please track https://github.com/pytorch/pytorch/issues/147851 "

149 "for the latest updates.")

150

151 - self.pointer_to_data: Dict[int, AllocationData] = {}

151 + self.pointer_to_data: dict[int, AllocationData] = {}

152 self.current_tag: str = CuMemAllocator.default_tag

153 - self.allocator_and_pools: Dict[str, Any] = {}

153 + self.allocator_and_pools: dict[str, Any] = {}

154

155 def python_malloc_callback(self, allocation_handle: HandleType) -> None:

156 """

@@ -172,7 +172,7 @@ class CuMemAllocator:

172

173 def sleep(

174 self,

175 - offload_tags: Optional[Union[Tuple[str, ...],

175 + offload_tags: Optional[Union[tuple[str, ...],

176 str]] = None) -> None:

177 """

178 Put the allocator in sleep mode.

vllm/distributed/communication_op.py

+2 -2

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Any, Dict, Optional, Union

3 +from typing import Any, Optional, Union

5 import torch

6 import torch.distributed

@@ -32,7 +32,7 @@ def tensor_model_parallel_gather(input_: torch.Tensor,

32 return get_tp_group().gather(input_, dst, dim)

35 -def broadcast_tensor_dict(tensor_dict: Optional[Dict[Any, Union[torch.Tensor,

35 +def broadcast_tensor_dict(tensor_dict: Optional[dict[Any, Union[torch.Tensor,

36 Any]]] = None,

37 src: int = 0):

38 if not torch.distributed.is_initialized():

vllm/distributed/device_communicators/base_device_communicator.py

+2 -2

@@ -1,5 +1,5 @@

1 # SPDX-License-Identifier: Apache-2.0

2 -from typing import Optional, Tuple

2 +from typing import Optional

4 import torch

5 import torch.distributed as dist

@@ -160,7 +160,7 @@ class DeviceCommunicatorBase:

160

161 def dispatch(

162 self, hidden_states: torch.Tensor,

163 - router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

163 + router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

164 """

165 Dispatch the hidden states and router logits to the appropriate device.

166 This is a no-op in the base class.

vllm/distributed/device_communicators/cpu_communicator.py

+2 -2

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import os

4 -from typing import List, Optional

4 +from typing import Optional

6 import torch

7 from torch.distributed import ProcessGroup

@@ -126,7 +126,7 @@ class _CPUSHMDistributed:

126

127 def gather(self,

128 input: torch.Tensor,

129 - gather_list: Optional[List[torch.Tensor]],

129 + gather_list: Optional[list[torch.Tensor]],

130 dst: int = -1,

131 group: Optional[ProcessGroup] = None) -> None:

132 # Note: different from the torch gather, here we use local dst rank.

vllm/distributed/device_communicators/cuda_communicator.py

+2 -2

@@ -1,6 +1,6 @@

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Optional, Tuple

3 +from typing import Optional

5 import torch

6 from torch.distributed import ProcessGroup

@@ -154,7 +154,7 @@ class CudaCommunicator(DeviceCommunicatorBase):

154

155 def dispatch(

156 self, hidden_states: torch.Tensor,

157 - router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

157 + router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

158 assert self.all2all_impl is not None

159 hidden_states, router_logits = self.all2all_impl.dispatch(

160 hidden_states, router_logits)

vllm/distributed/device_communicators/cuda_wrapper.py

+4 -4

@@ -6,7 +6,7 @@ convenient for use when we just need to call a few functions.

7 import ctypes

8 from dataclasses import dataclass

9 -from typing import Any, Dict, List, Optional

9 +from typing import Any, Optional

11 # this line makes it possible to directly load `libcudart.so` using `ctypes`

12 import torch # noqa

@@ -32,7 +32,7 @@ class cudaIpcMemHandle_t(ctypes.Structure):

32 class Function:

33 name: str

34 restype: Any

35 - argtypes: List[Any]

35 + argtypes: list[Any]

38 def find_loaded_library(lib_name) -> Optional[str]:

@@ -97,11 +97,11 @@ class CudaRTLibrary:

98 # class attribute to store the mapping from the path to the library

99 # to avoid loading the same library multiple times

100 - path_to_library_cache: Dict[str, Any] = {}

100 + path_to_library_cache: dict[str, Any] = {}

101

102 # class attribute to store the mapping from library path

103 # to the corresponding dictionary

104 - path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}

104 + path_to_dict_mapping: dict[str, dict[str, Any]] = {}

105

106 def __init__(self, so_file: Optional[str] = None):

107 if so_file is None:

vllm/distributed/device_communicators/custom_all_reduce.py

+4 -4

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from contextlib import contextmanager

4 -from typing import List, Optional, Union

4 +from typing import Optional, Union

6 import torch

7 import torch.distributed as dist

@@ -276,7 +276,7 @@ class CustomAllreduce:

276 @staticmethod

277 def create_shared_buffer(size_in_bytes: int,

278 group: Optional[ProcessGroup] = None,

279 - uncached: Optional[bool] = False) -> List[int]:

279 + uncached: Optional[bool] = False) -> list[int]:

280 pointer, handle = ops.allocate_shared_buffer_and_handle(size_in_bytes)

281

282 world_size = dist.get_world_size(group=group)

@@ -284,7 +284,7 @@ class CustomAllreduce:

284 handles = [None] * world_size

285 dist.all_gather_object(handles, handle, group=group)

286

287 - pointers: List[int] = []

287 + pointers: list[int] = []

288 for i, h in enumerate(handles):

289 if i == rank:

290 pointers.append(pointer) # type: ignore

@@ -293,7 +293,7 @@ class CustomAllreduce:

293 return pointers

294

295 @staticmethod

296 - def free_shared_buffer(pointers: List[int],

296 + def free_shared_buffer(pointers: list[int],

297 group: Optional[ProcessGroup] = None,

298 rank: Optional[int] = 0) -> None:

299 if rank is None:

vllm/distributed/device_communicators/custom_all_reduce_utils.py

+5 -4

@@ -7,8 +7,9 @@ import pickle

7 import subprocess

8 import sys

9 import tempfile

10 +from collections.abc import Sequence

11 from itertools import product

11 -from typing import Dict, List, Optional, Sequence

12 +from typing import Optional

14 import torch.distributed as dist

15 import torch.multiprocessing as mp

@@ -149,7 +150,7 @@ def can_actually_p2p(

150 p_src.join()

151 p_tgt.join()

152 assert p_src.exitcode == 0 and p_tgt.exitcode == 0

152 - result: List[bool] = []

153 + result: list[bool] = []

154 for src, tgt in zip(batch_src, batch_tgt):

155 a = result_queue.get()

156 b = result_queue.get()

@@ -175,7 +176,7 @@ def can_actually_p2p(

176 # e.g. used by different vllm engines. The device id in the cache file is a

177 # **local** device id, i.e. from 0 to num_dev-1, where num_dev is the number

178 # of visible devices in the vllm engine.

178 -_gpu_p2p_access_cache: Optional[Dict[str, bool]] = None

179 +_gpu_p2p_access_cache: Optional[dict[str, bool]] = None

180

181

182 def gpu_p2p_access_check(src: int, tgt: int) -> bool:

@@ -204,7 +205,7 @@ def gpu_p2p_access_check(src: int, tgt: int) -> bool:

205 # only the local master process (with local_rank == 0) can

206 # enter this block to calculate the cache

207 logger.info("generating GPU P2P access cache in %s", path)

207 - cache: Dict[str, bool] = {}

208 + cache: dict[str, bool] = {}

209 ids = list(range(num_dev))

210 # batch of all pairs of GPUs

211 batch_src, batch_tgt = zip(*list(product(ids, ids)))

vllm/distributed/device_communicators/pynccl_wrapper.py

+5 -5

@@ -24,7 +24,7 @@

24 import ctypes

25 import platform

26 from dataclasses import dataclass

27 -from typing import Any, Dict, List, Optional

27 +from typing import Any, Optional

29 import torch

30 from torch.distributed import ReduceOp

@@ -121,7 +121,7 @@ class ncclRedOpTypeEnum:

121 class Function:

122 name: str

123 restype: Any

124 - argtypes: List[Any]

124 + argtypes: list[Any]

125

126

127 class NCCLLibrary:

@@ -210,11 +210,11 @@ class NCCLLibrary:

210

211 # class attribute to store the mapping from the path to the library

212 # to avoid loading the same library multiple times

213 - path_to_library_cache: Dict[str, Any] = {}

213 + path_to_library_cache: dict[str, Any] = {}

214

215 # class attribute to store the mapping from library path

216 # to the corresponding dictionary

217 - path_to_dict_mapping: Dict[str, Dict[str, Any]] = {}

217 + path_to_dict_mapping: dict[str, dict[str, Any]] = {}

218

219 def __init__(self, so_file: Optional[str] = None):

220

@@ -238,7 +238,7 @@ class NCCLLibrary:

238 raise e

239

240 if so_file not in NCCLLibrary.path_to_dict_mapping:

241 - _funcs: Dict[str, Any] = {}

241 + _funcs: dict[str, Any] = {}

242 for func in NCCLLibrary.exported_functions:

243 f = getattr(self.lib, func.name)

244 f.restype = func.restype

vllm/distributed/device_communicators/shm_broadcast.py

+4 -4

@@ -8,7 +8,7 @@ from contextlib import contextmanager

8 from dataclasses import dataclass, field

9 from multiprocessing import shared_memory

10 from threading import Event

11 -from typing import Any, List, Optional, Tuple, Union

11 +from typing import Any, Optional, Union

12 from unittest.mock import patch

14 import torch

@@ -173,9 +173,9 @@ class ShmRingBuffer:

173

174 @dataclass

175 class Handle:

176 - local_reader_ranks: List[int] = field(default_factory=list)

176 + local_reader_ranks: list[int] = field(default_factory=list)

177

178 - buffer_handle: Optional[Tuple[int, int, int, str]] = None

178 + buffer_handle: Optional[tuple[int, int, int, str]] = None

179 local_subscribe_addr: Optional[str] = None

180 remote_subscribe_addr: Optional[str] = None

181 remote_addr_ipv6: bool = False

@@ -187,7 +187,7 @@ class MessageQueue:

187 self,

188 n_reader, # number of all readers

189 n_local_reader, # number of local readers through shared memory

190 - local_reader_ranks: Optional[List[int]] = None,

190 + local_reader_ranks: Optional[list[int]] = None,

191 max_chunk_bytes: int = 1024 * 1024 * 10,

192 max_chunks: int = 10,

193 connect_ip: Optional[str] = None,

vllm/distributed/kv_transfer/kv_connector/base.py

+6 -6

@@ -8,7 +8,7 @@ The class provides two primary abstract methods:

8 """

10 from abc import ABC, abstractmethod

11 -from typing import TYPE_CHECKING, List, Tuple, Union

11 +from typing import TYPE_CHECKING, Union

13 import torch

@@ -55,7 +55,7 @@ class KVConnectorBase(ABC):

55 self,

56 model_executable: torch.nn.Module,

57 model_input: "ModelInputForGPUWithSamplingMetadata",

58 - kv_caches: List[torch.Tensor],

58 + kv_caches: list[torch.Tensor],

59 hidden_or_intermediate_states: Union[torch.Tensor,

60 IntermediateTensors],

61 ) -> None:

@@ -71,7 +71,7 @@ class KVConnectorBase(ABC):

71 start and end layer information.

72 model_input (ModelInputForGPUWithSamplingMetadata): The input

73 metadata from vLLM.

74 - kv_caches (List[torch.Tensor]): List of KV caches (keys and values)

74 + kv_caches (list[torch.Tensor]): List of KV caches (keys and values)

75 for each layer.

76 hidden_or_intermediate_states (Union[torch.Tensor,

77 IntermediateTensors]):

@@ -88,8 +88,8 @@ class KVConnectorBase(ABC):

88 def recv_kv_caches_and_hidden_states(

89 self, model_executable: torch.nn.Module,

90 model_input: "ModelInputForGPUWithSamplingMetadata",

91 - kv_caches: List[torch.Tensor]

92 - ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,

91 + kv_caches: list[torch.Tensor]

92 + ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,

93 "ModelInputForGPUWithSamplingMetadata"]:

94 """

95 Receive KV caches and hidden states from the connector.

@@ -104,7 +104,7 @@ class KVConnectorBase(ABC):

104 The model executable from vLLM modelrunner.

105 model_input (ModelInputForGPUWithSamplingMetadata):

106 The model input from vLLM modelrunner.

107 - kv_caches (List[torch.Tensor]):

107 + kv_caches (list[torch.Tensor]):

108 List of KV caches for each layer.

109

110 Returns:

vllm/distributed/kv_transfer/kv_connector/factory.py

+3 -3

@@ -1,7 +1,7 @@

1 # SPDX-License-Identifier: Apache-2.0

3 import importlib

4 -from typing import TYPE_CHECKING, Callable, Dict, Type

4 +from typing import TYPE_CHECKING, Callable

6 import vllm.envs as envs

7 from vllm.distributed.kv_transfer.kv_connector.base import KVConnectorBaseType

@@ -18,7 +18,7 @@ logger = init_logger(__name__)

20 class KVConnectorFactory:

21 - _registry: Dict[str, Callable[[], Type[KVConnectorBaseType]]] = {}

21 + _registry: dict[str, Callable[[], type[KVConnectorBaseType]]] = {}

23 @classmethod

24 def register_connector(cls, name: str, module_path: str,

@@ -27,7 +27,7 @@ class KVConnectorFactory:

27 if name in cls._registry:

28 raise ValueError(f"Connector '{name}' is already registered.")

30 - def loader() -> Type[KVConnectorBaseType]:

30 + def loader() -> type[KVConnectorBaseType]:

31 module = importlib.import_module(module_path)

32 return getattr(module, class_name)

vllm/distributed/kv_transfer/kv_connector/lmcache_connector.py

+4 -4

@@ -7,7 +7,7 @@ The LMCacheConnector can (1) transfer KV caches between prefill vLLM worker

7 (2) offload and share KV caches.

8 """

10 -from typing import TYPE_CHECKING, List, Tuple, Union

10 +from typing import TYPE_CHECKING, Union

12 import torch

@@ -63,8 +63,8 @@ class LMCacheConnector(KVConnectorBase):

63 def recv_kv_caches_and_hidden_states(

64 self, model_executable: torch.nn.Module,

65 model_input: "ModelInputForGPUWithSamplingMetadata",

66 - kv_caches: List[torch.Tensor]

67 - ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,

66 + kv_caches: list[torch.Tensor]

67 + ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,

68 "ModelInputForGPUWithSamplingMetadata"]:

70 retrieve_status = self.lmcache_should_retrieve(model_input)

@@ -78,7 +78,7 @@ class LMCacheConnector(KVConnectorBase):

78 self,

79 model_executable: torch.nn.Module,

80 model_input: "ModelInputForGPUWithSamplingMetadata",

81 - kv_caches: List[torch.Tensor],

81 + kv_caches: list[torch.Tensor],

82 hidden_or_intermediate_states: Union[torch.Tensor,

83 IntermediateTensors],

84 ) -> None:

vllm/distributed/kv_transfer/kv_connector/mooncake_store_connector.py

+4 -4

@@ -6,7 +6,7 @@ The MooncakeStoreConnector transfers KV caches between prefill vLLM workers

6 database-style KVStore.

7 """

8 import hashlib

9 -from typing import TYPE_CHECKING, List, Tuple, Union

9 +from typing import TYPE_CHECKING, Union

11 import torch

@@ -70,7 +70,7 @@ class MooncakeStoreConnector(KVConnectorBase):

70 self,

71 model_executable: torch.nn.Module,

72 model_input: "ModelInputForGPUWithSamplingMetadata",

73 - kv_caches: List[torch.Tensor],

73 + kv_caches: list[torch.Tensor],

74 hidden_or_intermediate_states: Union[torch.Tensor,

75 IntermediateTensors],

76 ) -> None:

@@ -113,8 +113,8 @@ class MooncakeStoreConnector(KVConnectorBase):

113 def recv_kv_caches_and_hidden_states(

114 self, model_executable: torch.nn.Module,

115 model_input: "ModelInputForGPUWithSamplingMetadata",

116 - kv_caches: List[torch.Tensor]

117 - ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,

116 + kv_caches: list[torch.Tensor]

117 + ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,

118 "ModelInputForGPUWithSamplingMetadata"]:

119 bypass_model_exec = True

120 input_tokens_tensor = model_input.input_tokens

vllm/distributed/kv_transfer/kv_connector/simple_connector.py

+5 -5

@@ -8,7 +8,7 @@ MooncakePipe.

9 But the logic can be extended to support other pipe and lookup buffer.

10 """

11 -from typing import TYPE_CHECKING, List, Optional, Tuple, Union

11 +from typing import TYPE_CHECKING, Optional, Union

13 import torch

@@ -133,7 +133,7 @@ class SimpleConnector(KVConnectorBase):

133 )

134

135 def select(self, input_tokens: Optional[torch.Tensor],

136 - roi: Optional[torch.Tensor]) -> List[Optional[torch.Tensor]]:

136 + roi: Optional[torch.Tensor]) -> list[Optional[torch.Tensor]]:

137

138 assert self.consumer_buffer is not None, "Please initialize the "\

139 "consumer buffer before calling select."

@@ -152,7 +152,7 @@ class SimpleConnector(KVConnectorBase):

152 self,

153 model_executable: torch.nn.Module,

154 model_input: "ModelInputForGPUWithSamplingMetadata",

155 - kv_caches: List[torch.Tensor],

155 + kv_caches: list[torch.Tensor],

156 hidden_or_intermediate_states: Union[torch.Tensor,

157 IntermediateTensors],

158 ) -> None:

@@ -207,8 +207,8 @@ class SimpleConnector(KVConnectorBase):

207 def recv_kv_caches_and_hidden_states(

208 self, model_executable: torch.nn.Module,

209 model_input: "ModelInputForGPUWithSamplingMetadata",

210 - kv_caches: List[torch.Tensor]

211 - ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,

210 + kv_caches: list[torch.Tensor]

211 + ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,

212 "ModelInputForGPUWithSamplingMetadata"]:

213

214 # When bypass_model_exec is set to False, it means that at least for one

vllm/distributed/kv_transfer/kv_connector/v1/nixl_connector.py

+2 -2

@@ -5,13 +5,13 @@ import threading

5 import time

6 import uuid

7 from collections import defaultdict

8 +from collections.abc import Iterator

9 from dataclasses import dataclass

9 -from typing import TYPE_CHECKING, Any, Iterator

10 +from typing import TYPE_CHECKING, Any, Optional

12 import msgspec

13 import torch

14 import zmq

14 -from typing_extensions import Optional

16 from vllm import envs

17 from vllm.config import VllmConfig

vllm/distributed/kv_transfer/kv_connector_agent.py

+4 -4

@@ -5,7 +5,7 @@ This implementation is a shim wrapper on two APIs exposed by `kv_connector`:

5 1. `send_kv_caches_and_hidden_states`

6 2. `recv_kv_caches_and_hidden_states

7 """

8 -from typing import TYPE_CHECKING, List, Tuple, Union

8 +from typing import TYPE_CHECKING, Union

10 if TYPE_CHECKING:

11 from vllm.worker.model_runner import ModelInputForGPUWithSamplingMetadata

@@ -53,7 +53,7 @@ class KVTransferAgent:

53 self,

54 model_executable: torch.nn.Module,

55 model_input: "ModelInputForGPUWithSamplingMetadata",

56 - kv_caches: List[torch.Tensor],

56 + kv_caches: list[torch.Tensor],

57 hidden_or_intermediate_states: Union[torch.Tensor,

58 IntermediateTensors],

59 ) -> None:

@@ -68,8 +68,8 @@ class KVTransferAgent:

68 def recv_kv_caches_and_hidden_states(

69 self, model_executable: torch.nn.Module,

70 model_input: "ModelInputForGPUWithSamplingMetadata",

71 - kv_caches: List[torch.Tensor]

72 - ) -> Tuple[Union[torch.Tensor, IntermediateTensors], bool,

71 + kv_caches: list[torch.Tensor]

72 + ) -> tuple[Union[torch.Tensor, IntermediateTensors], bool,

73 "ModelInputForGPUWithSamplingMetadata"]:

75 return self.connector.recv_kv_caches_and_hidden_states(

vllm/distributed/kv_transfer/kv_lookup_buffer/base.py

+3 -3

@@ -13,7 +13,7 @@ These classes above are abstracted behind class `KVCacheBufferBase`.

13 """

15 from abc import ABC, abstractmethod

16 -from typing import List, Optional

16 +from typing import Optional

18 import torch

@@ -93,7 +93,7 @@ class KVLookupBufferBase(KVCacheBufferBase):

93 @abstractmethod

94 def drop_select(

95 self, input_tokens: Optional[torch.Tensor],

96 - roi: Optional[torch.Tensor]) -> List[Optional[torch.Tensor]]:

96 + roi: Optional[torch.Tensor]) -> list[Optional[torch.Tensor]]:

97 """Select and *drop* KV cache entries from the lookup buffer.

99 The functionality is similar to the following python statements

@@ -111,7 +111,7 @@ class KVLookupBufferBase(KVCacheBufferBase):

111 roi (torch.Tensor): A binary mask on top of the input tokens

112

113 Returns:

114 - List[Optional[torch.Tensor]]: A list of tensors. Can be None.

114 + list[Optional[torch.Tensor]]: A list of tensors. Can be None.

115

116 Raises:

117 NotImplementedError: This method must be implemented in subclasses.

vllm/distributed/kv_transfer/kv_lookup_buffer/simple_buffer.py

+7 -7

@@ -11,7 +11,7 @@

11 """

12 import threading

13 from collections import deque

14 -from typing import Deque, List, Optional, Union

14 +from typing import Optional, Union

16 import torch

@@ -38,7 +38,7 @@ class SimpleBuffer(KVLookupBufferBase):

38 data_pipe: on device (e.g. GPU)

39 """

41 - self.buffer: Deque[List[torch.Tensor]] = deque()

41 + self.buffer: deque[list[torch.Tensor]] = deque()

43 self.buffer_size = 0

44 self.buffer_size_threshold = buffer_size_thresh

@@ -50,8 +50,8 @@ class SimpleBuffer(KVLookupBufferBase):

50 self.normal_signal = torch.tensor([0], device="cpu")

51 self.end_signal = None

53 - def _matches(self, tokens_roi_sender: List[torch.Tensor],

54 - tokens_roi_recver: List[torch.Tensor]):

53 + def _matches(self, tokens_roi_sender: list[torch.Tensor],

54 + tokens_roi_recver: list[torch.Tensor]):

56 # tokens_roi_sender: tokens and roi of the producer (in the buffer)

57 # tokens_roi_recver: tokens and roi of the consumer (query)

@@ -88,7 +88,7 @@ class SimpleBuffer(KVLookupBufferBase):

88 tensor = tensor.float()

89 self.data_pipe.send_tensor(tensor)

91 - def _get_element_size(self, data: Optional[Union[List, torch.Tensor]]):

91 + def _get_element_size(self, data: Optional[Union[list, torch.Tensor]]):

93 if isinstance(data, torch.Tensor):

94 return data.element_size() * data.numel()

@@ -151,7 +151,7 @@ class SimpleBuffer(KVLookupBufferBase):

151 tokens_roi_recver = [input_tokens, roi]

152

153 def is_buffer_available(

154 - tokens_roi_recver: List[torch.Tensor], ) -> bool:

154 + tokens_roi_recver: list[torch.Tensor], ) -> bool:

155 # perform input tokens and roi matching

156 # FIXME: this matching is O(n), ideally it should be O(1)

157 # but this buffer size won't (and shouldn't) be too large so

@@ -184,7 +184,7 @@ class SimpleBuffer(KVLookupBufferBase):

184

185 def drop_select(

186 self, input_tokens: Optional[torch.Tensor],

187 - roi: Optional[torch.Tensor]) -> List[Optional[torch.Tensor]]:

187 + roi: Optional[torch.Tensor]) -> list[Optional[torch.Tensor]]:

188

189 assert self.request_handling_thread is None, \

190 "drop_select should be called by the KV cache consumer "\

vllm/distributed/kv_transfer/kv_pipe/pynccl_pipe.py

+3 -3

@@ -15,7 +15,7 @@

15 import threading

16 import time

17 from concurrent.futures import ThreadPoolExecutor

18 -from typing import Callable, Dict, Optional, Tuple

18 +from typing import Callable, Optional

20 import torch

@@ -35,7 +35,7 @@ class BrokenPipeException(Exception):

35 super().__init__(self.message)

38 -Metadata = Dict[str, Optional[torch.Tensor]]

38 +Metadata = dict[str, Optional[torch.Tensor]]

41 class PyNcclPipe(KVPipeBase):

@@ -83,7 +83,7 @@ class PyNcclPipe(KVPipeBase):

84 def _get_device_send_recv_impl(

85 self, group: StatelessProcessGroup

86 - ) -> Tuple[Callable[[torch.Tensor, int], None], Callable[

86 + ) -> tuple[Callable[[torch.Tensor, int], None], Callable[

87 [torch.Tensor, int], None]]:

89 send: Callable[[torch.Tensor, int], None]

vllm/distributed/parallel_state.py

+22 -22

@@ -29,7 +29,7 @@ from collections import namedtuple

29 from contextlib import contextmanager, nullcontext

30 from dataclasses import dataclass

31 from multiprocessing import shared_memory

32 -from typing import Any, Callable, Dict, List, Optional, Tuple, Union

32 +from typing import Any, Callable, Optional, Union

33 from unittest.mock import patch

35 import torch

@@ -54,15 +54,15 @@ TensorMetadata = namedtuple("TensorMetadata", ["device", "dtype", "size"])

56 def _split_tensor_dict(

57 - tensor_dict: Dict[str, Union[torch.Tensor, Any]]

58 -) -> Tuple[List[Tuple[str, Any]], List[torch.Tensor]]:

57 + tensor_dict: dict[str, Union[torch.Tensor, Any]]

58 +) -> tuple[list[tuple[str, Any]], list[torch.Tensor]]:

59 """Split the tensor dictionary into two parts:

60 1. A list of (key, value) pairs. If the value is a tensor, it is replaced

61 by its metadata.

62 2. A list of tensors.

63 """

64 - metadata_list: List[Tuple[str, Any]] = []

65 - tensor_list: List[torch.Tensor] = []

64 + metadata_list: list[tuple[str, Any]] = []

65 + tensor_list: list[torch.Tensor] = []

66 for key, value in tensor_dict.items():

67 if isinstance(value, torch.Tensor):

68 # Note: we cannot use `value.device` here,

@@ -78,7 +78,7 @@ def _split_tensor_dict(

78 return metadata_list, tensor_list

81 -_group_name_counter: Dict[str, int] = {}

81 +_group_name_counter: dict[str, int] = {}

84 def _get_unique_name(name: str) -> str:

@@ -94,7 +94,7 @@ def _get_unique_name(name: str) -> str:

94 return newname

97 -_groups: Dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}

97 +_groups: dict[str, Callable[[], Optional["GroupCoordinator"]]] = {}

100 def _register_group(group: "GroupCoordinator") -> None:

@@ -182,7 +182,7 @@ class GroupCoordinator:

182

183 # available attributes:

184 rank: int # global rank

185 - ranks: List[int] # global ranks in the group

185 + ranks: list[int] # global ranks in the group

186 world_size: int # size of the group

187 # difference between `local_rank` and `rank_in_group`:

188 # if we have a group of size 4 across two nodes:

@@ -201,7 +201,7 @@ class GroupCoordinator:

201

202 def __init__(

203 self,

204 - group_ranks: List[List[int]],

204 + group_ranks: list[list[int]],

205 local_rank: int,

206 torch_distributed_backend: Union[str, Backend],

207 use_device_communicator: bool,

@@ -435,7 +435,7 @@ class GroupCoordinator:

435 return recv[0]

436

437 def broadcast_object_list(self,

438 - obj_list: List[Any],

438 + obj_list: list[Any],

439 src: int = 0,

440 group: Optional[ProcessGroup] = None):

441 """Broadcast the input object list.

@@ -518,11 +518,11 @@ class GroupCoordinator:

518

519 def broadcast_tensor_dict(

520 self,

521 - tensor_dict: Optional[Dict[str, Union[torch.Tensor, Any]]] = None,

521 + tensor_dict: Optional[dict[str, Union[torch.Tensor, Any]]] = None,

522 src: int = 0,

523 group: Optional[ProcessGroup] = None,

524 metadata_group: Optional[ProcessGroup] = None

525 - ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:

525 + ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

526 """Broadcast the input tensor dictionary.

527 NOTE: `src` is the local rank of the source rank.

528 """

@@ -536,7 +536,7 @@ class GroupCoordinator:

536

537 rank_in_group = self.rank_in_group

538 if rank_in_group == src:

539 - metadata_list: List[Tuple[Any, Any]] = []

539 + metadata_list: list[tuple[Any, Any]] = []

540 assert isinstance(

541 tensor_dict,

542 dict), (f"Expecting a dictionary, got {type(tensor_dict)}")

@@ -603,10 +603,10 @@ class GroupCoordinator:

603

604 def send_tensor_dict(

605 self,

606 - tensor_dict: Dict[str, Union[torch.Tensor, Any]],

606 + tensor_dict: dict[str, Union[torch.Tensor, Any]],

607 dst: Optional[int] = None,

608 all_gather_group: Optional["GroupCoordinator"] = None,

609 - ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:

609 + ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

610 """Send the input tensor dictionary.

611 NOTE: `dst` is the local rank of the source rank.

612 """

@@ -626,7 +626,7 @@ class GroupCoordinator:

626 dst = (self.rank_in_group + 1) % self.world_size

627 assert dst < self.world_size, f"Invalid dst rank ({dst})"

628

629 - metadata_list: List[Tuple[Any, Any]] = []

629 + metadata_list: list[tuple[Any, Any]] = []

630 assert isinstance(

631 tensor_dict,

632 dict), f"Expecting a dictionary, got {type(tensor_dict)}"

@@ -661,7 +661,7 @@ class GroupCoordinator:

661 self,

662 src: Optional[int] = None,

663 all_gather_group: Optional["GroupCoordinator"] = None,

664 - ) -> Optional[Dict[str, Union[torch.Tensor, Any]]]:

664 + ) -> Optional[dict[str, Union[torch.Tensor, Any]]]:

665 """Recv the input tensor dictionary.

666 NOTE: `src` is the local rank of the source rank.

667 """

@@ -682,7 +682,7 @@ class GroupCoordinator:

682 assert src < self.world_size, f"Invalid src rank ({src})"

683

684 recv_metadata_list = self.recv_object(src=src)

685 - tensor_dict: Dict[str, Any] = {}

685 + tensor_dict: dict[str, Any] = {}

686 for key, value in recv_metadata_list:

687 if isinstance(value, TensorMetadata):

688 tensor = torch.empty(value.size,

@@ -764,7 +764,7 @@ class GroupCoordinator:

764

765 def dispatch(

766 self, hidden_states: torch.Tensor,

767 - router_logits: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:

767 + router_logits: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:

768 if self.device_communicator is not None:

769 return self.device_communicator.dispatch(hidden_states,

770 router_logits)

@@ -782,7 +782,7 @@ def get_world_group() -> GroupCoordinator:

782 return _WORLD

783

784

785 -def init_world_group(ranks: List[int], local_rank: int,

785 +def init_world_group(ranks: list[int], local_rank: int,

786 backend: str) -> GroupCoordinator:

787 return GroupCoordinator(

788 group_ranks=[ranks],

@@ -794,7 +794,7 @@ def init_world_group(ranks: List[int], local_rank: int,

794

795

796 def init_model_parallel_group(

797 - group_ranks: List[List[int]],

797 + group_ranks: list[list[int]],

798 local_rank: int,

799 backend: str,

800 use_message_queue_broadcaster: bool = False,

@@ -1182,7 +1182,7 @@ def cleanup_dist_env_and_memory(shutdown_ray: bool = False):

1182

1183

1184 def in_the_same_node_as(pg: Union[ProcessGroup, StatelessProcessGroup],

1185 - source_rank: int = 0) -> List[bool]:

1185 + source_rank: int = 0) -> list[bool]:

1186 """

1187 This is a collective operation that returns if each rank is in the same node

1188 as the source rank. It tests if processes are attached to the same

vllm/distributed/utils.py

+7 -6

@@ -10,7 +10,8 @@ import pickle

10 import socket

11 import time

12 from collections import deque

13 -from typing import Any, Deque, Dict, Optional, Sequence, Tuple

13 +from collections.abc import Sequence

14 +from typing import Any, Optional

16 import torch

17 from torch.distributed import ProcessGroup, TCPStore

@@ -69,7 +70,7 @@ def split_tensor_along_last_dim(

72 def get_pp_indices(num_hidden_layers: int, pp_rank: int,

72 - pp_size: int) -> Tuple[int, int]:

73 + pp_size: int) -> tuple[int, int]:

74 """Try to evenly distribute layers across partitions.

76 If the number of layers is not divisible by the number of partitions,

@@ -132,15 +133,15 @@ class StatelessProcessGroup:

133 data_expiration_seconds: int = 3600 # 1 hour

134

135 # dst rank -> counter

135 - send_dst_counter: Dict[int, int] = dataclasses.field(default_factory=dict)

136 + send_dst_counter: dict[int, int] = dataclasses.field(default_factory=dict)

137 # src rank -> counter

137 - recv_src_counter: Dict[int, int] = dataclasses.field(default_factory=dict)

138 + recv_src_counter: dict[int, int] = dataclasses.field(default_factory=dict)

139 broadcast_send_counter: int = 0

139 - broadcast_recv_src_counter: Dict[int, int] = dataclasses.field(

140 + broadcast_recv_src_counter: dict[int, int] = dataclasses.field(

141 default_factory=dict)

142

143 # A deque to store the data entries, with key and timestamp.

143 - entries: Deque[Tuple[str,

144 + entries: deque[tuple[str,

145 float]] = dataclasses.field(default_factory=deque)

146

147 def __post_init__(self):