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

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 deprecated type hinting in `vllm/compilation` (#18072)

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

Harry Mellor committed May 13, 2025 at 16:32 UTC 19324d660c61a63c6ea3dfbb18995d255c05ee6d

13 files changed +70 -69

pyproject.toml

-1

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

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

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

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

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

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

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

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

vllm/compilation/backends.py

+17 -16

@@ -5,8 +5,9 @@ import dataclasses

5 import os

6 import pprint

7 import time

8 +from collections.abc import Sequence

9 from contextlib import ExitStack

9 -from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple

10 +from typing import Any, Callable, Optional

11 from unittest.mock import patch

13 import torch

@@ -56,7 +57,7 @@ class CompilerManager:

57 """

59 def __init__(self, compilation_config: CompilationConfig):

59 - self.cache: Dict[Tuple[Optional[int], int, str], Any] = dict()

60 + self.cache: dict[tuple[Optional[int], int, str], Any] = dict()

61 self.is_cache_updated = False

62 self.compilation_config = compilation_config

63 self.compiler = make_compiler(compilation_config)

@@ -90,7 +91,7 @@ class CompilerManager:

92 def load(self,

93 graph: fx.GraphModule,

93 - example_inputs: List[Any],

94 + example_inputs: list[Any],

95 graph_index: int,

96 runtime_shape: Optional[int] = None) -> Optional[Callable]:

97 if (runtime_shape, graph_index, self.compiler.name) not in self.cache:

@@ -186,7 +187,7 @@ class SplitItem:

187

188

189 def split_graph(graph: fx.GraphModule,

189 - ops: List[str]) -> Tuple[fx.GraphModule, List[SplitItem]]:

190 + ops: list[str]) -> tuple[fx.GraphModule, list[SplitItem]]:

191 # split graph by ops

192 subgraph_id = 0

193 node_to_subgraph_id = {}

@@ -252,7 +253,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):

253 """

254

255 def __init__(self, module: torch.fx.GraphModule,

255 - compile_submod_names: List[str], vllm_config: VllmConfig,

256 + compile_submod_names: list[str], vllm_config: VllmConfig,

257 graph_pool, vllm_backend: "VllmBackend"):

258 super().__init__(module)

259 from torch._guards import detect_fake_mode

@@ -274,8 +275,8 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):

275 return super().run(*fake_args)

276

277 def call_module(self, target: torch.fx.node.Target,

277 - args: Tuple[torch.fx.node.Argument,

278 - ...], kwargs: Dict[str, Any]) -> Any:

278 + args: tuple[torch.fx.node.Argument,

279 + ...], kwargs: dict[str, Any]) -> Any:

280 assert isinstance(target, str)

281 output = super().call_module(target, args, kwargs)

282

@@ -326,12 +327,12 @@ class VllmBackend:

327 graph: fx.GraphModule

328 # the stiching graph module for all the piecewise graphs

329 split_gm: fx.GraphModule

329 - piecewise_graphs: List[SplitItem]

330 + piecewise_graphs: list[SplitItem]

331 returned_callable: Callable

332 # Inductor passes to run on the graph pre-defunctionalization

333 post_grad_passes: Sequence[Callable]

333 - sym_tensor_indices: List[int]

334 - input_buffers: List[torch.Tensor]

334 + sym_tensor_indices: list[int]

335 + input_buffers: list[torch.Tensor]

336 compiler_manager: CompilerManager

337

338 def __init__(

@@ -573,14 +574,14 @@ class ConcreteSizeEntry:

574

575 # for cudagraph debugging, track the input addresses

576 # during capture, and check if they are the same during replay

576 - input_addresses: Optional[List[int]] = None

577 + input_addresses: Optional[list[int]] = None

578

579

580 class PiecewiseBackend:

581

582 def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,

583 graph_pool: Any, piecewise_compile_index: int,

583 - total_piecewise_compiles: int, sym_shape_indices: List[int],

584 + total_piecewise_compiles: int, sym_shape_indices: list[int],

585 compiled_graph_for_general_shape: Callable,

586 vllm_backend: VllmBackend):

587 """

@@ -608,9 +609,9 @@ class PiecewiseBackend:

609 self.is_last_graph = (

610 piecewise_compile_index == total_piecewise_compiles - 1)

611

611 - self.compile_sizes: Set[int] = set(

612 + self.compile_sizes: set[int] = set(

613 self.compilation_config.compile_sizes)

613 - self.cudagraph_capture_sizes: Set[int] = set(

614 + self.cudagraph_capture_sizes: set[int] = set(

615 self.compilation_config.cudagraph_capture_sizes

616 ) if self.compilation_config.use_cudagraph else set()

617

@@ -624,11 +625,11 @@ class PiecewiseBackend:

625

626 # the entries for different shapes that we need to either

627 # compile or capture cudagraph

627 - self.concrete_size_entries: Dict[int, ConcreteSizeEntry] = {}

628 + self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {}

629

630 # to_be_compiled_sizes tracks the remaining sizes to compile,

631 # and updates during the compilation process, so we need to copy it

631 - self.to_be_compiled_sizes: Set[int] = self.compile_sizes.copy()

632 + self.to_be_compiled_sizes: set[int] = self.compile_sizes.copy()

633 for shape in self.compile_sizes.union(self.cudagraph_capture_sizes):

634 self.concrete_size_entries[shape] = ConcreteSizeEntry(

635 runtime_shape=shape,

vllm/compilation/compiler_interface.py

+19 -19

@@ -4,7 +4,7 @@ import copy

4 import hashlib

5 import os

6 from contextlib import ExitStack

7 -from typing import Any, Callable, Dict, List, Optional, Tuple

7 +from typing import Any, Callable, Optional

8 from unittest.mock import patch

10 import torch

@@ -48,11 +48,11 @@ class CompilerInterface:

48 def compile(

49 self,

50 graph: fx.GraphModule,

51 - example_inputs: List[Any],

52 - compiler_config: Dict[str, Any],

51 + example_inputs: list[Any],

52 + compiler_config: dict[str, Any],

53 runtime_shape: Optional[int] = None,

54 key: Optional[str] = None,

55 - ) -> Tuple[Optional[Callable], Optional[Any]]:

55 + ) -> tuple[Optional[Callable], Optional[Any]]:

56 """

57 Compile the graph with the given example inputs and compiler config,

58 with a runtime shape. If the `runtime_shape` is None, it means

@@ -82,7 +82,7 @@ class CompilerInterface:

82 def load(self,

83 handle: Any,

84 graph: fx.GraphModule,

85 - example_inputs: List[Any],

85 + example_inputs: list[Any],

86 graph_index: int,

87 runtime_shape: Optional[int] = None) -> Callable:

88 """

@@ -120,7 +120,7 @@ class AlwaysHitShapeEnv:

120 """

121

122 def __init__(self) -> None:

123 - self.guards: List[Any] = []

123 + self.guards: list[Any] = []

124

125 def evaluate_guards_expression(self, *args, **kwargs):

126 return True

@@ -132,8 +132,8 @@ class AlwaysHitShapeEnv:

132 return ""

133

134

135 -def get_inductor_factors() -> List[Any]:

136 - factors: List[Any] = []

135 +def get_inductor_factors() -> list[Any]:

136 + factors: list[Any] = []

137 # summarize system state

138 from torch._inductor.codecache import CacheBase

139 system_factors = CacheBase.get_system()

@@ -169,11 +169,11 @@ class InductorStandaloneAdaptor(CompilerInterface):

169 def compile(

170 self,

171 graph: fx.GraphModule,

172 - example_inputs: List[Any],

173 - compiler_config: Dict[str, Any],

172 + example_inputs: list[Any],

173 + compiler_config: dict[str, Any],

174 runtime_shape: Optional[int] = None,

175 key: Optional[str] = None,

176 - ) -> Tuple[Optional[Callable], Optional[Any]]:

176 + ) -> tuple[Optional[Callable], Optional[Any]]:

177 current_config = {}

178 if compiler_config is not None:

179 current_config.update(compiler_config)

@@ -201,7 +201,7 @@ class InductorStandaloneAdaptor(CompilerInterface):

201 def load(self,

202 handle: Any,

203 graph: fx.GraphModule,

204 - example_inputs: List[Any],

204 + example_inputs: list[Any],

205 graph_index: int,

206 runtime_shape: Optional[int] = None) -> Callable:

207 assert isinstance(handle, tuple)

@@ -256,11 +256,11 @@ class InductorAdaptor(CompilerInterface):

256 def compile(

257 self,

258 graph: fx.GraphModule,

259 - example_inputs: List[Any],

260 - compiler_config: Dict[str, Any],

259 + example_inputs: list[Any],

260 + compiler_config: dict[str, Any],

261 runtime_shape: Optional[int] = None,

262 key: Optional[str] = None,

263 - ) -> Tuple[Optional[Callable], Optional[Any]]:

263 + ) -> tuple[Optional[Callable], Optional[Any]]:

264 from torch._inductor.compile_fx import compile_fx

265 current_config = {}

266 if compiler_config is not None:

@@ -420,7 +420,7 @@ class InductorAdaptor(CompilerInterface):

420 def load(self,

421 handle: Any,

422 graph: fx.GraphModule,

423 - example_inputs: List[Any],

423 + example_inputs: list[Any],

424 graph_index: int,

425 runtime_shape: Optional[int] = None) -> Callable:

426 assert isinstance(handle, tuple)

@@ -522,11 +522,11 @@ class EagerAdaptor(CompilerInterface):

522 def compile(

523 self,

524 graph: fx.GraphModule,

525 - example_inputs: List[Any],

526 - compiler_config: Dict[str, Any],

525 + example_inputs: list[Any],

526 + compiler_config: dict[str, Any],

527 runtime_shape: Optional[int] = None,

528 key: Optional[str] = None,

529 - ) -> Tuple[Optional[Callable], Optional[Any]]:

529 + ) -> tuple[Optional[Callable], Optional[Any]]:

530 # we don't need to compile the graph, just return the graph itself.

531 # It does not support caching, return None for the handle.

532 return graph, None

vllm/compilation/decorators.py

+4 -4

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

1 # SPDX-License-Identifier: Apache-2.0

3 import inspect

4 -from typing import Callable, Dict, List, Optional, TypeVar, Union, overload

4 +from typing import Callable, Optional, TypeVar, Union, overload

5 from unittest.mock import patch

7 import torch

@@ -25,7 +25,7 @@ _T = TypeVar("_T", bound=type[nn.Module])

25 @overload

26 def support_torch_compile(

27 *,

28 - dynamic_arg_dims: Optional[Dict[str, Union[int, List[int]]]],

28 + dynamic_arg_dims: Optional[dict[str, Union[int, list[int]]]],

29 ) -> Callable[[_T], _T]:

30 ...

@@ -38,7 +38,7 @@ def support_torch_compile(cls: _T) -> _T:

38 def support_torch_compile(

39 cls: Optional[_T] = None,

40 *,

41 - dynamic_arg_dims: Optional[Dict[str, Union[int, List[int]]]] = None,

41 + dynamic_arg_dims: Optional[dict[str, Union[int, list[int]]]] = None,

42 ) -> Union[Callable[[_T], _T], _T]:

43 """

44 A decorator to add support for compiling the forward method of a class.

@@ -131,7 +131,7 @@ def support_torch_compile(

131

132 def _support_torch_compile(

133 cls: _T,

134 - dynamic_arg_dims: Dict[str, Union[int, List[int]]],

134 + dynamic_arg_dims: dict[str, Union[int, list[int]]],

135 ) -> _T:

136 """

137 A decorator to add support for compiling the forward method of a class.

vllm/compilation/fix_functionalization.py

+8 -7

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

1 # SPDX-License-Identifier: Apache-2.0

3 import operator

4 -from typing import Dict, Iterable, List, Optional, Tuple, Union

4 +from collections.abc import Iterable

5 +from typing import Optional, Union

7 import torch

8 from torch._higher_order_ops.auto_functionalize import auto_functionalized

@@ -27,7 +28,7 @@ class FixFunctionalizationPass(VllmInductorPass):

28 self.begin()

29 self.dump_graph(graph, "before_fix_functionalization")

30 - self.nodes_to_remove: List[torch.fx.Node] = []

31 + self.nodes_to_remove: list[torch.fx.Node] = []

32 count = 0

33 for node in graph.nodes:

34 if not is_func(node, auto_functionalized):

@@ -117,8 +118,8 @@ class FixFunctionalizationPass(VllmInductorPass):

118 def defunctionalize(self,

119 graph: torch.fx.Graph,

120 node: torch.fx.Node,

120 - mutated_args: Dict[int, Union[torch.fx.Node, str]],

121 - args: Optional[Tuple[Union[torch.fx.Node, str],

121 + mutated_args: dict[int, Union[torch.fx.Node, str]],

122 + args: Optional[tuple[Union[torch.fx.Node, str],

123 ...]] = None):

124 """

125 De-functionalize a node by replacing it with a call to the original.

@@ -130,7 +131,7 @@ class FixFunctionalizationPass(VllmInductorPass):

131 self._remove(node)

132

133 def replace_users_with_mutated_args(self, node: torch.fx.Node,

133 - mutated_args: Dict[int,

134 + mutated_args: dict[int,

135 Union[torch.fx.Node,

136 str]]):

137 """

@@ -146,7 +147,7 @@ class FixFunctionalizationPass(VllmInductorPass):

147 user.replace_all_uses_with(arg)

148 self._remove(user)

149

149 - def getitem_users(self, node: torch.fx.Node) -> Dict[int, torch.fx.Node]:

150 + def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]:

151 """

152 Returns the operator.getitem users of the auto-functionalized node,

153 indexed by the index they are getting.

@@ -161,7 +162,7 @@ class FixFunctionalizationPass(VllmInductorPass):

162 def insert_defunctionalized(self,

163 graph: torch.fx.Graph,

164 node: torch.fx.Node,

164 - args: Optional[Tuple[Union[torch.fx.Node, str],

165 + args: Optional[tuple[Union[torch.fx.Node, str],

166 ...]] = None):

167 """

168 Insert a new defunctionalized node into the graph before node.

vllm/compilation/fusion.py

+5 -5

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

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Callable, Dict, List, NamedTuple, Optional, Tuple

3 +from typing import Callable, NamedTuple, Optional

5 import torch

6 import torch._inductor.pattern_matcher as pm

@@ -57,7 +57,7 @@ kFp8StaticTensorSym = QuantKey(FP8_DTYPE, True, True, True)

57 kFp8DynamicTensorSym = QuantKey(FP8_DTYPE, False, True, True)

58 kFp8DynamicTokenSym = QuantKey(FP8_DTYPE, False, False, True)

60 -QUANT_OPS: Dict[QuantKey, OpOverload] = {

60 +QUANT_OPS: dict[QuantKey, OpOverload] = {

61 kFp8StaticTensorSym: torch.ops._C.static_scaled_fp8_quant.default, # noqa

62 kFp8DynamicTensorSym:

63 torch.ops._C.dynamic_scaled_fp8_quant.default, # noqa

@@ -80,7 +80,7 @@ class FusedRMSQuantKey(NamedTuple):

80 f"{'' if self.fused_add else 'out'} residual)")

83 -FUSED_OPS: Dict[FusedRMSQuantKey, OpOverload] = {

83 +FUSED_OPS: dict[FusedRMSQuantKey, OpOverload] = {

84 FusedRMSQuantKey(kFp8StaticTensorSym, False):

85 torch.ops._C.rms_norm_static_fp8_quant.default, # noqa

86 FusedRMSQuantKey(kFp8StaticTensorSym, True):

@@ -101,7 +101,7 @@ class QuantMultiOutputMatch(MultiOutputMatch):

101 self.QUANT_OP = quant_op # in-place quant op

102 self.FUSED_OP = fused_op # in-place fused quant op

103

104 - def insert_fused_node(self, fused_return_mapping: Dict[int, Tuple[fx.Node,

104 + def insert_fused_node(self, fused_return_mapping: dict[int, tuple[fx.Node,

105 int]],

106 **kwargs):

107 """

@@ -548,7 +548,7 @@ class FusionPass(VllmInductorPass):

548 "FusionPass singleton instance already exists"

549 super().__init__(config)

550

551 - self.matches: List[MultiOutputMatch] = []

551 + self.matches: list[MultiOutputMatch] = []

552 self.patterns: PatternMatcherPass = PatternMatcherPass(

553 pass_name="fusion_pass")

554

vllm/compilation/fx_utils.py

+2 -1

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

1 # SPDX-License-Identifier: Apache-2.0

3 import operator

4 -from typing import Iterable, Optional

4 +from collections.abc import Iterable

5 +from typing import Optional

7 from torch import fx

8 from torch._higher_order_ops.auto_functionalize import auto_functionalized

vllm/compilation/inductor_pass.py

+2 -2

@@ -5,7 +5,7 @@ import inspect

5 import json

6 import types

7 from contextlib import contextmanager

8 -from typing import Any, Callable, Dict, Optional, Union

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

10 import torch

11 from torch import fx

@@ -83,7 +83,7 @@ class InductorPass(CustomGraphPass):

83 return hasher.hexdigest()

85 @staticmethod

86 - def hash_dict(dict_: Dict[Any, Any]):

86 + def hash_dict(dict_: dict[Any, Any]):

87 """

88 Utility method to hash a dictionary, can alternatively be used for uuid.

89 :return: A sha256 hash of the json rep of the dictionary.

vllm/compilation/multi_output_match.py

+3 -3

@@ -3,7 +3,7 @@

3 import abc

4 import operator

5 from abc import abstractmethod

6 -from typing import Iterable, List, Tuple

6 +from collections.abc import Iterable

8 from torch import fx

9 from torch._higher_order_ops.auto_functionalize import auto_functionalized

@@ -56,7 +56,7 @@ class MultiOutputMatch(abc.ABC):

56 raise NotImplementedError

58 @property

59 - def nodes(self) -> List[fx.Node]:

59 + def nodes(self) -> list[fx.Node]:

60 return self.match.nodes

62 @property

@@ -87,7 +87,7 @@ class MultiOutputMatch(abc.ABC):

87 return self.graph.inserting_after(last_node_in_match)

89 def insert_getitems(self, tuple_node: fx.Node,

90 - indices: Iterable[int]) -> Tuple[fx.Node, ...]:

90 + indices: Iterable[int]) -> tuple[fx.Node, ...]:

91 """

92 Insert operator.getitem nodes to extract elements from a tuple node.

vllm/compilation/noop_elimination.py

+2 -1

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

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import Iterable, Union

3 +from collections.abc import Iterable

4 +from typing import Union

6 import torch.fx

7 from torch import SymInt

vllm/compilation/pass_manager.py

+1 -3

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

1 # SPDX-License-Identifier: Apache-2.0

3 -from typing import List

4 -

3 from torch import fx as fx

5 from vllm.config import VllmConfig

@@ -34,7 +32,7 @@ class PostGradPassManager(CustomGraphPass):

32 """

34 def __init__(self):

37 - self.passes: List[VllmInductorPass] = []

35 + self.passes: list[VllmInductorPass] = []

37 def __call__(self, graph: fx.Graph):

38 shape = get_pass_context().runtime_shape

vllm/compilation/sequence_parallelism.py

+5 -5

@@ -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._inductor.pattern_matcher as pm

@@ -125,7 +125,7 @@ class MiddleAllReduceRMSNormPattern(AllReduceRMSNormPattern):

125 residual: torch.Tensor,

126 mm_1: torch.Tensor,

127 rms_norm_weights: torch.Tensor,

128 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

129 all_reduce = tensor_model_parallel_all_reduce(mm_1)

130

131 rmsnorm = torch.ops.higher_order.auto_functionalized(

@@ -142,7 +142,7 @@ class MiddleAllReduceRMSNormPattern(AllReduceRMSNormPattern):

142 residual: torch.Tensor,

143 mm_1: torch.Tensor,

144 rms_norm_weights: torch.Tensor,

145 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

146 tp = get_tp_group()

147 tp_size = get_tensor_model_parallel_world_size()

148 reduce_scatter = torch.ops.vllm.reduce_scatter.default(

@@ -190,7 +190,7 @@ class LastAllReduceRMSNormPattern(AllReduceRMSNormPattern):

190 residual: torch.Tensor,

191 mm_1: torch.Tensor,

192 rms_norm_weights: torch.Tensor,

193 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

194 all_reduce = tensor_model_parallel_all_reduce(mm_1)

195

196 rmsnorm = torch.ops.higher_order.auto_functionalized(

@@ -207,7 +207,7 @@ class LastAllReduceRMSNormPattern(AllReduceRMSNormPattern):

207 residual: torch.Tensor,

208 mm_1: torch.Tensor,

209 rms_norm_weights: torch.Tensor,

210 - ) -> Tuple[torch.Tensor, torch.Tensor]:

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

211 tp = get_tp_group()

212 tp_size = get_tensor_model_parallel_world_size()

213 reduce_scatter = torch.ops.vllm.reduce_scatter.default(

vllm/compilation/wrapper.py

+2 -2

@@ -5,7 +5,7 @@ import sys

5 from abc import abstractmethod

6 from contextlib import contextmanager

7 from types import CodeType

8 -from typing import Callable, List, Optional

8 +from typing import Callable, Optional

10 import torch

@@ -48,7 +48,7 @@ class TorchCompileWrapperWithCustomDispatcher:

49 self.compiled_callable = compiled_callable

50 self.original_code_object = self.__class__.forward.__code__

51 - self.compiled_codes: List[CodeType] = []

51 + self.compiled_codes: list[CodeType] = []

52 torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook)

54 # read the env var to determine whether to use the custom dispatcher