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

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, 11:00 UTC

Update some more deprecated type hinting (#17998)

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

Harry Mellor committed May 13, 2025 at 00:49 UTC 9d7ea9dbbf570098dba44e4f00948bef2510fc8d

10 files changed +73 -73

pyproject.toml

+3 -1

@@ -79,7 +79,9 @@ exclude = [

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

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

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

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

82 +"vllm/model_executor/layers/**/*.py" = ["UP006", "UP035"]

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

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

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

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

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

vllm/model_executor/custom_op.py

+1 -3

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

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Dict, Type

4 -

3 import torch.nn as nn

5 from vllm.config import get_current_vllm_config

@@ -138,7 +136,7 @@ class CustomOp(nn.Module):

136 # Examples:

137 # - MyOp.enabled()

138 # - op_registry["my_op"].enabled()

141 - op_registry: Dict[str, Type['CustomOp']] = {}

139 + op_registry: dict[str, type['CustomOp']] = {}

140

141 # Decorator to register custom ops.

142 @classmethod

vllm/model_executor/guided_decoding/guidance_logits_processors.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 import os

3 -from typing import Any, List

3 +from typing import Any

5 import llguidance

6 import llguidance.hf

@@ -62,7 +62,7 @@ class GuidanceLogitsProcessor:

63 def __call__(

64 self,

65 - input_ids: List[int],

65 + input_ids: list[int],

66 scores: torch.Tensor,

67 ) -> torch.Tensor:

68 # we initialize the guidance model here

vllm/model_executor/guided_decoding/guided_fields.py

+5 -5

@@ -1,16 +1,16 @@

1 # SPDX-License-Identifier: Apache-2.0

3 from dataclasses import dataclass

4 -from typing import Dict, List, Optional, TypedDict, Union

4 +from typing import Optional, TypedDict, Union

6 from pydantic import BaseModel

9 # These classes are deprecated, see SamplingParams

10 class LLMGuidedOptions(TypedDict, total=False):

11 - guided_json: Union[Dict, BaseModel, str]

11 + guided_json: Union[dict, BaseModel, str]

12 guided_regex: str

13 - guided_choice: List[str]

13 + guided_choice: list[str]

14 guided_grammar: str

15 guided_decoding_backend: str

16 guided_whitespace_pattern: str

@@ -20,9 +20,9 @@ class LLMGuidedOptions(TypedDict, total=False):

20 @dataclass

21 class GuidedDecodingRequest:

22 """One of the fields will be used to retrieve the logit processor."""

23 - guided_json: Optional[Union[Dict, BaseModel, str]] = None

23 + guided_json: Optional[Union[dict, BaseModel, str]] = None

24 guided_regex: Optional[str] = None

25 - guided_choice: Optional[List[str]] = None

25 + guided_choice: Optional[list[str]] = None

26 guided_grammar: Optional[str] = None

27 guided_decoding_backend: Optional[str] = None

28 guided_whitespace_pattern: Optional[str] = None

vllm/model_executor/guided_decoding/outlines_decoding.py

+2 -2

@@ -6,7 +6,7 @@ import os

6 from enum import Enum

7 from json import dumps as json_dumps

8 from re import escape as regex_escape

9 -from typing import Optional, Tuple, Union

9 +from typing import Optional, Union

11 from transformers import PreTrainedTokenizerBase

@@ -111,7 +111,7 @@ def get_local_outlines_guided_decoding_logits_processor(

111

112 def _get_guide_and_mode(

113 guided_params: GuidedDecodingParams

114 -) -> Union[Tuple[str, GuidedDecodingMode], Tuple[None, None]]:

114 +) -> Union[tuple[str, GuidedDecodingMode], tuple[None, None]]:

115 if guided_params.json:

116 if isinstance(guided_params.json, dict):

117 # turn dict into hashable string

vllm/model_executor/guided_decoding/outlines_logits_processors.py

+8 -8

@@ -19,7 +19,7 @@ import copy

19 import json

20 from collections import defaultdict

21 from functools import lru_cache

22 -from typing import Callable, DefaultDict, Dict, List, Optional, Union

22 +from typing import Callable, Optional, Union

24 import numpy as np

25 import torch

@@ -53,10 +53,10 @@ class BaseLogitsProcessor:

53 self._guide: Guide = guide

54 self._reasoner: Optional[ReasoningParser] = reasoner

55 # CFGState is used for the FSM state for CFGGuide

56 - self._fsm_state: DefaultDict[int, Union[int,

56 + self._fsm_state: defaultdict[int, Union[int,

57 CFGState]] = defaultdict(int)

59 - def __call__(self, input_ids: List[int],

59 + def __call__(self, input_ids: list[int],

60 scores: torch.Tensor) -> torch.Tensor:

61 """Use the FSM to bias the logits before sampling the next token."""

@@ -160,7 +160,7 @@ class RegexLogitsProcessor(BaseLogitsProcessor):

160

161 class JSONLogitsProcessor(RegexLogitsProcessor):

162

163 - def __init__(self, schema: Union[str, Dict, BaseModel],

163 + def __init__(self, schema: Union[str, dict, BaseModel],

164 tokenizer: PreTrainedTokenizerBase,

165 whitespace_pattern: Union[str, None],

166 reasoner: Optional[ReasoningParser]):

@@ -181,7 +181,7 @@ class JSONLogitsProcessor(RegexLogitsProcessor):

181 """

182 if isinstance(schema, type(BaseModel)):

183 schema_str = json.dumps(schema.model_json_schema())

184 - elif isinstance(schema, Dict):

184 + elif isinstance(schema, dict):

185 schema_str = json.dumps(schema)

186 elif isinstance(schema, str):

187 schema_str = schema

@@ -252,11 +252,11 @@ def _adapt_tokenizer(tokenizer: PreTrainedTokenizerBase):

252 return string

253

254 def change_decoder(

255 - decoder: Callable[[List[int]],

256 - str]) -> Callable[[List[int]], List[str]]:

255 + decoder: Callable[[list[int]],

256 + str]) -> Callable[[list[int]], list[str]]:

257 """Sync vLLM's decoder with the outlines by returning list."""

258

259 - def new_decoder(inp_tokens: List[int]) -> List[str]:

259 + def new_decoder(inp_tokens: list[int]) -> list[str]:

260 if (isinstance(inp_tokens, list) and len(inp_tokens) == 1

261 and isinstance(inp_tokens[0], list)):

262 inp_tokens = inp_tokens[0]

vllm/model_executor/guided_decoding/xgrammar_decoding.py

+2 -2

@@ -6,7 +6,7 @@ from __future__ import annotations

6 import json

7 import re

8 from dataclasses import dataclass, field

9 -from typing import TYPE_CHECKING, Any, List

9 +from typing import TYPE_CHECKING, Any

11 import torch

@@ -273,7 +273,7 @@ class GrammarConfig:

273 return re.sub(r'(["\\])', r'\\\1', s)

274

275 @staticmethod

276 - def choice_as_grammar(choice: List[str] | None) -> str:

276 + def choice_as_grammar(choice: list[str] | None) -> str:

277 if choice is None:

278 raise ValueError("Choice is not set")

279 escaped_choices = (GrammarConfig.escape_ebnf_string(c) for c in choice)

vllm/model_executor/pooling_metadata.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

3 from dataclasses import dataclass

4 -from typing import Any, Dict, List, Tuple

4 +from typing import Any

6 import torch

@@ -23,9 +23,9 @@ class PoolingMetadata:

24 def __init__(

25 self,

26 - seq_groups: List[Tuple[List[int], PoolingParams]],

27 - seq_data: Dict[int, Any], # Specific data related to sequences

28 - prompt_lens: List[int],

26 + seq_groups: list[tuple[list[int], PoolingParams]],

27 + seq_data: dict[int, Any], # Specific data related to sequences

28 + prompt_lens: list[int],

29 ) -> None:

30 self.seq_groups = seq_groups

31 self.seq_data = seq_data

vllm/model_executor/sampling_metadata.py

+44 -44

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

3 from array import array

4 from dataclasses import dataclass

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

5 +from typing import Optional

7 import torch

@@ -25,10 +25,10 @@ class SequenceGroupToSample:

25 # |-- query_len ---|

27 # Sequence ids for the sequence group in a previous step.

28 - seq_ids: List[int]

28 + seq_ids: list[int]

29 sampling_params: SamplingParams

30 # seq_id -> sequence data.

31 - seq_data: Dict[int, SequenceData]

31 + seq_data: dict[int, SequenceData]

32 # The length of the sequence (all tokens seen in the past + new token to

33 # compute attention) of the sequence group. None if it is in a decode

34 # stage.

@@ -44,9 +44,9 @@ class SequenceGroupToSample:

44 is_prompt: bool

45 # Query token indices from logits. to compute prompt logprob. Empty if

46 # prompt logprob is not required.

47 - prompt_logprob_indices: List[int]

47 + prompt_logprob_indices: list[int]

48 # Sample token indices from logits. Empty if sampling is not required.

49 - sample_indices: List[int]

49 + sample_indices: list[int]

51 @property

52 def do_sample(self):

@@ -78,7 +78,7 @@ class SamplingMetadataCache:

78 """Used to cache SamplingMetadata objects between scheduler iterations"""

80 def __init__(self):

81 - self._seq_group_to_sample_cache: Dict[int, PyObjectCache] = {}

81 + self._seq_group_to_sample_cache: dict[int, PyObjectCache] = {}

83 def get_cached_seq_group_to_sample(self, num_seqs):

84 if num_seqs not in self._seq_group_to_sample_cache:

@@ -130,9 +130,9 @@ class SamplingMetadata:

130

131 def __init__(

132 self,

133 - seq_groups: List[SequenceGroupToSample],

133 + seq_groups: list[SequenceGroupToSample],

134 selected_token_indices: torch.Tensor,

135 - categorized_sample_indices: Dict[SamplingType, torch.Tensor],

135 + categorized_sample_indices: dict[SamplingType, torch.Tensor],

136 num_prompts: int,

137 skip_sampler_cpu_output: bool = False,

138 reuse_sampling_tensors: bool = False,

@@ -146,12 +146,12 @@ class SamplingMetadata:

146

147 @staticmethod

148 def prepare(

149 - seq_group_metadata_list: List[SequenceGroupMetadata],

150 - seq_lens: List[int],

151 - query_lens: List[int],

149 + seq_group_metadata_list: list[SequenceGroupMetadata],

150 + seq_lens: list[int],

151 + query_lens: list[int],

152 device: str,

153 pin_memory: bool,

154 - generators: Optional[Dict[str, torch.Generator]] = None,

154 + generators: Optional[dict[str, torch.Generator]] = None,

155 cache: Optional[SamplingMetadataCache] = None,

156 ) -> "SamplingMetadata":

157 (

@@ -195,16 +195,16 @@ class SamplingMetadata:

195

196

197 def _prepare_seq_groups(

198 - seq_group_metadata_list: List[SequenceGroupMetadata],

199 - seq_lens: List[int],

200 - query_lens: List[int],

198 + seq_group_metadata_list: list[SequenceGroupMetadata],

199 + seq_lens: list[int],

200 + query_lens: list[int],

201 device: str,

202 - generators: Optional[Dict[str, torch.Generator]] = None,

202 + generators: Optional[dict[str, torch.Generator]] = None,

203 cache: Optional[SamplingMetadataCache] = None,

204 -) -> Tuple[

205 - List[SequenceGroupToSample],

206 - List[int],

207 - Dict[SamplingType, List[int]],

204 +) -> tuple[

205 + list[SequenceGroupToSample],

206 + list[int],

207 + dict[SamplingType, list[int]],

208 int,

209 ]:

210 """Prepare sequence groups and indices for sampling.

@@ -227,17 +227,17 @@ def _prepare_seq_groups(

227 num_prompts: Total number of prompts from `seq_group_metadata_list`.

228 """

229 # Batched sequence groups for the current model forward stsep.

230 - seq_groups: List[SequenceGroupToSample] = []

230 + seq_groups: list[SequenceGroupToSample] = []

231 # A list of token indices to sample/compute logprob. It is used to

232 # prune the outcome logits from the model for the performance.

233 - selected_token_indices: List[int] = []

233 + selected_token_indices: list[int] = []

234 # Used for selected_token_indices.

235 model_output_idx = 0

236

237 # Sampling type -> (

238 # indices to sample/prompt logprob within pruned output logits,

239 # indices to sample within pruned logits)

240 - categorized_sample_indices: Dict[SamplingType, List[int]] = {

240 + categorized_sample_indices: dict[SamplingType, list[int]] = {

241 t: []

242 for t in SamplingType

243 }

@@ -265,9 +265,9 @@ def _prepare_seq_groups(

265 # If the current seq group is in decode stage, it is None.

266 seq_len: Optional[int] = None

267 query_len: Optional[int] = None

268 - prompt_logprob_indices: List[int] = (sample_obj.prompt_logprob_indices

268 + prompt_logprob_indices: list[int] = (sample_obj.prompt_logprob_indices

269 if cache is not None else [])

270 - sample_indices: List[int] = (sample_obj.sample_indices

270 + sample_indices: list[int] = (sample_obj.sample_indices

271 if cache is not None else [])

272 do_sample = seq_group_metadata.do_sample

273

@@ -389,16 +389,16 @@ class SamplingTensors:

389 vocab_size: int,

390 device: torch.device,

391 dtype: torch.dtype,

392 - ) -> Tuple["SamplingTensors", bool, bool, bool]:

393 - prompt_tokens: List[array] = []

394 - output_tokens: List[array] = []

395 - top_ks: List[int] = []

396 - temperatures: List[float] = []

397 - top_ps: List[float] = []

398 - min_ps: List[float] = []

399 - presence_penalties: List[float] = []

400 - frequency_penalties: List[float] = []

401 - repetition_penalties: List[float] = []

392 + ) -> tuple["SamplingTensors", bool, bool, bool]:

393 + prompt_tokens: list[array] = []

394 + output_tokens: list[array] = []

395 + top_ks: list[int] = []

396 + temperatures: list[float] = []

397 + top_ps: list[float] = []

398 + min_ps: list[float] = []

399 + presence_penalties: list[float] = []

400 + frequency_penalties: list[float] = []

401 + repetition_penalties: list[float] = []

402 do_penalties = False

403 do_top_p_top_k = False

404 do_min_p = False

@@ -496,15 +496,15 @@ class SamplingTensors:

496 @classmethod

497 def from_lists(

498 cls,

499 - temperatures: List[float],

500 - top_ps: List[float],

501 - top_ks: List[int],

502 - min_ps: List[float],

503 - presence_penalties: List[float],

504 - frequency_penalties: List[float],

505 - repetition_penalties: List[float],

506 - prompt_tokens: List[array],

507 - output_tokens: List[array],

499 + temperatures: list[float],

500 + top_ps: list[float],

501 + top_ks: list[int],

502 + min_ps: list[float],

503 + presence_penalties: list[float],

504 + frequency_penalties: list[float],

505 + repetition_penalties: list[float],

506 + prompt_tokens: list[array],

507 + output_tokens: list[array],

508 vocab_size: int,

509 device: torch.device,

510 dtype: torch.dtype,

vllm/model_executor/utils.py

+2 -2

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

1 # SPDX-License-Identifier: Apache-2.0

2 """Utils for model executor."""

3 -from typing import Any, Dict, Optional

3 +from typing import Any, Optional

5 import torch

@@ -12,7 +12,7 @@ def set_random_seed(seed: int) -> None:

13 def set_weight_attrs(

14 weight: torch.Tensor,

15 - weight_attrs: Optional[Dict[str, Any]],

15 + weight_attrs: Optional[dict[str, Any]],

16 ):

17 """Set attributes on a weight tensor.