diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index 0c7bbfe599b0..826d1744d88a 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -141,14 +141,14 @@ def produce_guards_expression(self, *args, **kwargs): return "" -def wrap_inductor(graph, +def wrap_inductor(graph: fx.GraphModule, example_inputs, additional_inductor_config, compilation_config: CompilationConfig, graph_index: int = 0, num_graphs: int = 1, runtime_shape: Optional[int] = None, - use_inductor: bool = True): + use_inductor: bool = True) -> Any: if graph_index == 0: # before compiling the first graph, record the start time global compilation_start_time @@ -209,7 +209,7 @@ def wrap_inductor(graph, returns_tuple = graph_returns_tuple(graph) # this is the graph we return to Dynamo to run - def compiled_graph(*args): + def compiled_graph(*args) -> Optional[fx.CompiledFxGraph]: # convert args to list list_args = list(args) graph_output = inductor_compiled_graph(list_args) @@ -247,7 +247,7 @@ def _check_can_cache(*args, **kwargs): # see https://github.com/pytorch/pytorch/blob/9f5ebf3fc609105a74eab4ccc24932d6353ff566/torch/_inductor/codecache.py#L1221 # noqa return - def _get_shape_env(): + def _get_shape_env() -> AlwaysHitShapeEnv: return AlwaysHitShapeEnv() with patch(# for hijacking the hash of the compiled graph @@ -537,7 +537,7 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable: example_inputs[x].clone() for x in self.sym_tensor_indices ] - def copy_and_call(*args): + def copy_and_call(*args) -> fx.GraphModule: list_args = list(args) for i, index in enumerate(self.sym_tensor_indices): runtime_tensor = list_args[index] diff --git a/vllm/compilation/multi_output_match.py b/vllm/compilation/multi_output_match.py index 0ad648abfbb3..b6bcecdc89e2 100644 --- a/vllm/compilation/multi_output_match.py +++ b/vllm/compilation/multi_output_match.py @@ -7,6 +7,7 @@ from torch._higher_order_ops.auto_functionalize import auto_functionalized from torch._inductor import pattern_matcher as pm from torch._ops import OpOverload +from torch.fx import Node from vllm.compilation.fx_utils import find_auto_fn @@ -97,7 +98,7 @@ def insert_getitems(self, tuple_node: fx.Node, self.graph.call_function(operator.getitem, (tuple_node, idx)) for idx in indices) - def insert_auto_fn(self, op: OpOverload, kwargs): + def insert_auto_fn(self, op: OpOverload, kwargs) -> Node: """ Insert an auto_functionalized node with the given op and kwargs. """ diff --git a/vllm/compilation/pass_manager.py b/vllm/compilation/pass_manager.py index fb522ae053e9..34f5f355798b 100644 --- a/vllm/compilation/pass_manager.py +++ b/vllm/compilation/pass_manager.py @@ -1,4 +1,4 @@ -from typing import List +from typing import Any, Dict, List from torch import fx as fx @@ -53,7 +53,7 @@ def add(self, pass_: InductorPass): assert isinstance(pass_, InductorPass) self.passes.append(pass_) - def __getstate__(self): + def __getstate__(self) -> Dict[str, List[Any]]: """ Custom pickling for the pass manager, as some passes cannot be pickled. Pickling occurs because the pass manager is set as the value of