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

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 `platform`, `plugins`, `triton_utils`, `vllm_flash_attn` (#18129)

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

Harry Mellor committed May 14, 2025 at 13:28 UTC c8ea982d9b86e145a16092017528d068a7f94630

6 files changed +18 -24

pyproject.toml

-5

@@ -78,13 +78,8 @@ exclude = [

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

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

80 "vllm/model_executor/models/**/*.py" = ["UP006", "UP035"]

81 -"vllm/platforms/**/*.py" = ["UP006", "UP035"]

82 -"vllm/plugins/**/*.py" = ["UP006", "UP035"]

81 "vllm/prompt_adapter/**/*.py" = ["UP006", "UP035"]

82 "vllm/spec_decode/**/*.py" = ["UP006", "UP035"]

85 -"vllm/transformers_utils/**/*.py" = ["UP006", "UP035"]

86 -"vllm/triton_utils/**/*.py" = ["UP006", "UP035"]

87 -"vllm/vllm_flash_attn/**/*.py" = ["UP006", "UP035"]

83 "vllm/worker/**/*.py" = ["UP006", "UP035"]

84 "vllm/utils.py" = ["UP006", "UP035"]

vllm/platforms/cuda.py

+6 -7

@@ -5,8 +5,7 @@ pynvml. However, it should not initialize cuda context.

6 import os

7 from functools import wraps

8 -from typing import (TYPE_CHECKING, Callable, List, Optional, Tuple, TypeVar,

9 - Union)

8 +from typing import TYPE_CHECKING, Callable, Optional, TypeVar, Union

10 import torch

11 from typing_extensions import ParamSpec

@@ -56,7 +55,7 @@ class CudaPlatformBase(Platform):

55 device_control_env_var: str = "CUDA_VISIBLE_DEVICES"

57 @property

59 - def supported_dtypes(self) -> List[torch.dtype]:

58 + def supported_dtypes(self) -> list[torch.dtype]:

59 if self.has_device_capability(80):

60 # Ampere and Hopper or later NVIDIA GPUs.

61 return [torch.bfloat16, torch.float16, torch.float32]

@@ -93,7 +92,7 @@ class CudaPlatformBase(Platform):

92 return True

94 @classmethod

96 - def is_fully_connected(cls, device_ids: List[int]) -> bool:

95 + def is_fully_connected(cls, device_ids: list[int]) -> bool:

96 raise NotImplementedError

98 @classmethod

@@ -335,7 +334,7 @@ class NvmlCudaPlatform(CudaPlatformBase):

334 @with_nvml_context

335 def has_device_capability(

336 cls,

338 - capability: Union[Tuple[int, int], int],

337 + capability: Union[tuple[int, int], int],

338 device_id: int = 0,

339 ) -> bool:

340 try:

@@ -365,7 +364,7 @@ class NvmlCudaPlatform(CudaPlatformBase):

364

365 @classmethod

366 @with_nvml_context

368 - def is_fully_connected(cls, physical_device_ids: List[int]) -> bool:

367 + def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:

368 """

369 query if the set of gpus are fully connected by nvlink (1 hop)

370 """

@@ -430,7 +429,7 @@ class NonNvmlCudaPlatform(CudaPlatformBase):

429 return device_props.total_memory

430

431 @classmethod

433 - def is_fully_connected(cls, physical_device_ids: List[int]) -> bool:

432 + def is_fully_connected(cls, physical_device_ids: list[int]) -> bool:

433 logger.exception(

434 "NVLink detection not possible, as context support was"

435 " not found. Assuming no NVLink available.")

vllm/platforms/interface.py

+3 -3

@@ -4,7 +4,7 @@ import os

4 import platform

5 import random

6 from platform import uname

7 -from typing import TYPE_CHECKING, NamedTuple, Optional, Tuple, Union

7 +from typing import TYPE_CHECKING, NamedTuple, Optional, Union

9 import numpy as np

10 import torch

@@ -200,7 +200,7 @@ class Platform:

200 @classmethod

201 def has_device_capability(

202 cls,

203 - capability: Union[Tuple[int, int], int],

203 + capability: Union[tuple[int, int], int],

204 device_id: int = 0,

205 ) -> bool:

206 """

@@ -362,7 +362,7 @@ class Platform:

362 raise NotImplementedError

363

364 @classmethod

365 - def get_infinity_values(cls, dtype: torch.dtype) -> Tuple[float, float]:

365 + def get_infinity_values(cls, dtype: torch.dtype) -> tuple[float, float]:

366 """

367 Return the platform specific values for (-inf, inf)

368 """

vllm/platforms/rocm.py

+5 -5

@@ -2,7 +2,7 @@

3 import os

4 from functools import cache, lru_cache, wraps

5 -from typing import TYPE_CHECKING, Dict, List, Optional

5 +from typing import TYPE_CHECKING, Optional

7 import torch

@@ -35,7 +35,7 @@ except ImportError as e:

35 logger.warning("Failed to import from vllm._rocm_C with %r", e)

37 # Models not supported by ROCm.

38 -_ROCM_UNSUPPORTED_MODELS: List[str] = []

38 +_ROCM_UNSUPPORTED_MODELS: list[str] = []

40 # Models partially supported by ROCm.

41 # Architecture -> Reason.

@@ -43,7 +43,7 @@ _ROCM_SWA_REASON = ("Sliding window attention (SWA) is not yet supported in "

43 "Triton flash attention. For half-precision SWA support, "

44 "please use CK flash attention by setting "

45 "`VLLM_USE_TRITON_FLASH_ATTN=0`")

46 -_ROCM_PARTIALLY_SUPPORTED_MODELS: Dict[str, str] = {

46 +_ROCM_PARTIALLY_SUPPORTED_MODELS: dict[str, str] = {

47 "Qwen2ForCausalLM":

48 _ROCM_SWA_REASON,

49 "MistralForCausalLM":

@@ -58,7 +58,7 @@ _ROCM_PARTIALLY_SUPPORTED_MODELS: Dict[str, str] = {

58 "excessive use of shared memory. If this happens, disable Triton FA "

59 "by setting `VLLM_USE_TRITON_FLASH_ATTN=0`")

60 }

61 -_ROCM_DEVICE_ID_NAME_MAP: Dict[str, str] = {

61 +_ROCM_DEVICE_ID_NAME_MAP: dict[str, str] = {

62 "0x74a0": "AMD_Instinct_MI300A",

63 "0x74a1": "AMD_Instinct_MI300X",

64 "0x74b5": "AMD_Instinct_MI300X", # MI300X VF

@@ -203,7 +203,7 @@ class RocmPlatform(Platform):

203

204 @staticmethod

205 @with_amdsmi_context

206 - def is_fully_connected(physical_device_ids: List[int]) -> bool:

206 + def is_fully_connected(physical_device_ids: list[int]) -> bool:

207 """

208 Query if the set of gpus are fully connected by xgmi (1 hop)

209 """

vllm/platforms/tpu.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import TYPE_CHECKING, Optional, Tuple, Union, cast

3 +from typing import TYPE_CHECKING, Optional, Union, cast

5 import torch

6 from tpu_info import device

@@ -73,7 +73,7 @@ class TpuPlatform(Platform):

73 return "vllm.lora.punica_wrapper.punica_tpu.PunicaWrapperTPU"

75 @classmethod

76 - def get_infinity_values(cls, dtype: torch.dtype) -> Tuple[float, float]:

76 + def get_infinity_values(cls, dtype: torch.dtype) -> tuple[float, float]:

77 return torch.finfo(dtype).min, torch.finfo(dtype).max

79 @classmethod

vllm/plugins/__init__.py

+2 -2

@@ -2,7 +2,7 @@

3 import logging

4 import os

5 -from typing import Callable, Dict

5 +from typing import Callable

7 import torch

@@ -14,7 +14,7 @@ logger = logging.getLogger(__name__)

14 plugins_loaded = False

17 -def load_plugins_by_group(group: str) -> Dict[str, Callable]:

17 +def load_plugins_by_group(group: str) -> dict[str, Callable]:

18 import sys

19 if sys.version_info < (3, 10):

20 from importlib_metadata import entry_points