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

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

[Misc][BE] Type coverage for vllm/compilation [2/3] (#31744)

Lucas Kabela committed Jan 9, 2026 at 15:30 UTC aaf4b70aae5976115eb5f8caead3d4eeb3657136

12 files changed +161 -91

vllm/compilation/backends.py

+6 -6

@@ -179,7 +179,7 @@ class CompilerManager:

179 example_inputs: list[Any],

180 graph_index: int,

181 compile_range: Range,

182 - ) -> Callable | None:

182 + ) -> Callable[..., Any] | None:

183 if (compile_range, graph_index, self.compiler.name) not in self.cache:

184 return None

185 handle = self.cache[(compile_range, graph_index, self.compiler.name)]

@@ -199,7 +199,7 @@ class CompilerManager:

199 self,

200 graph: fx.GraphModule,

201 example_inputs: list[Any],

202 - additional_inductor_config,

202 + additional_inductor_config: dict[str, Any],

203 compilation_config: CompilationConfig,

204 compile_range: Range,

205 graph_index: int = 0,

@@ -355,7 +355,7 @@ def split_graph(

355 compilation_start_time = 0.0

356

357

358 -class PiecewiseCompileInterpreter(torch.fx.Interpreter):

358 +class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]

359 """Code adapted from `torch.fx.passes.shape_prop.ShapeProp`.

360 It runs the given graph with fake inputs, and compile some

361 submodules specified by `compile_submod_names` with the given

@@ -506,9 +506,9 @@ class VllmBackend:

506 # the stiching graph module for all the piecewise graphs

507 split_gm: fx.GraphModule

508 piecewise_graphs: list[SplitItem]

509 - returned_callable: Callable

509 + returned_callable: Callable[..., Any]

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

511 - post_grad_passes: Sequence[Callable]

511 + post_grad_passes: Sequence[Callable[..., Any]]

512 sym_tensor_indices: list[int]

513 input_buffers: list[torch.Tensor]

514 compiler_manager: CompilerManager

@@ -821,7 +821,7 @@ class VllmBackend:

821 ]

822

823 # this is the callable we return to Dynamo to run

824 - def copy_and_call(*args):

824 + def copy_and_call(*args: Any) -> Any:

825 list_args = list(args)

826 for i, index in enumerate(self.sym_tensor_indices):

827 runtime_tensor = list_args[index]

vllm/compilation/caching.py

+16 -7

@@ -4,6 +4,8 @@

4 import inspect

5 import os

6 import pickle

7 +from collections.abc import Callable, Sequence

8 +from typing import Any, Literal

9 from unittest.mock import patch

11 import torch

@@ -25,7 +27,7 @@ assert isinstance(SerializableCallable, type)

27 logger = init_logger(__name__)

28 -class VllmSerializableFunction(SerializableCallable):

30 +class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]

31 """

32 A wrapper around a compiled function by vllm. It will forward the tensor

33 inputs to the compiled function and return the result.

@@ -38,8 +40,13 @@ class VllmSerializableFunction(SerializableCallable):

40 """

42 def __init__(

41 - self, graph_module, example_inputs, prefix, optimized_call, is_encoder=False

42 - ):

43 + self,

44 + graph_module: torch.fx.GraphModule,

45 + example_inputs: Sequence[Any],

46 + prefix: str,

47 + optimized_call: Callable[..., Any],

48 + is_encoder: bool = False,

49 + ) -> None:

50 assert isinstance(graph_module, torch.fx.GraphModule)

51 self.graph_module = graph_module

52 self.example_inputs = example_inputs

@@ -53,7 +60,7 @@ class VllmSerializableFunction(SerializableCallable):

60 if sym_input is not None:

61 self.shape_env = sym_input.node.shape_env

56 - def __call__(self, *args, **kwargs):

63 + def __call__(self, *args: Any, **kwargs: Any) -> Any:

64 return self.optimized_call(*args, **kwargs)

66 @classmethod

@@ -73,7 +80,9 @@ class VllmSerializableFunction(SerializableCallable):

81 graph_reducer_override = GraphPickler.reducer_override

76 - def _graph_reducer_override(self, obj):

83 + def _graph_reducer_override(

84 + self: GraphPickler, obj: Any

85 + ) -> tuple[Callable[..., Any], tuple[Any, ...]] | Any:

86 if (

87 inspect.isclass(obj)

88 and issubclass(obj, sympy.Function)

@@ -114,7 +123,7 @@ class VllmSerializableFunction(SerializableCallable):

123 get_current_vllm_config(), state["prefix"], is_encoder

124 )

125

117 - def optimized_call(*example_inputs):

126 + def optimized_call(*example_inputs: Any) -> Any:

127 """

128 On the first run of the optimized call, we rerun the compiler

129 backend which should result in a cache hit. After the backend

@@ -136,7 +145,7 @@ class VllmSerializableFunction(SerializableCallable):

145 return fn

146

147 @property

139 - def co_name(self):

148 + def co_name(self) -> Literal["VllmSerializableFunction"]:

149 """

150 Used for depyf debugging.

151 """

vllm/compilation/cuda_graph.py

+12 -10

@@ -42,7 +42,9 @@ class CUDAGraphLogging:

42 "Count",

43 ]

45 - def __init__(self, cg_mode: CUDAGraphMode, cg_capture_sizes: list[int] | None):

45 + def __init__(

46 + self, cg_mode: CUDAGraphMode, cg_capture_sizes: list[int] | None

47 + ) -> None:

48 self.reset()

49 self.cg_mode = str(cg_mode)

50 self.cg_capture_sizes = str(cg_capture_sizes or [])

@@ -54,10 +56,10 @@ class CUDAGraphLogging:

56 "**CUDAGraph Stats:**\n\n"

57 )

57 - def reset(self):

58 - self.stats = []

59 + def reset(self) -> None:

60 + self.stats: list[CUDAGraphStat] = []

60 - def observe(self, cudagraph_stat: CUDAGraphStat):

62 + def observe(self, cudagraph_stat: CUDAGraphStat) -> None:

63 self.stats.append(cudagraph_stat)

65 def generate_metric_table(self) -> str:

@@ -109,7 +111,7 @@ class CUDAGraphLogging:

111 + "\n"

112 )

113

112 - def log(self, log_fn=logger.info):

114 + def log(self, log_fn: Callable[..., Any] = logger.info) -> None:

115 if not self.stats:

116 return

117 log_fn(self.generate_metric_table())

@@ -161,11 +163,11 @@ class CUDAGraphWrapper:

163

164 def __init__(

165 self,

164 - runnable: Callable,

166 + runnable: Callable[..., Any],

167 vllm_config: VllmConfig,

168 runtime_mode: CUDAGraphMode,

169 cudagraph_options: CUDAGraphOptions | None = None,

168 - ):

170 + ) -> None:

171 self.runnable = runnable

172 self.vllm_config = vllm_config

173 self.runtime_mode = runtime_mode

@@ -189,7 +191,7 @@ class CUDAGraphWrapper:

191 # cudagraphs for.

192 self.concrete_cudagraph_entries: dict[BatchDescriptor, CUDAGraphEntry] = {}

193

192 - def __getattr__(self, key: str):

194 + def __getattr__(self, key: str) -> Any:

195 # allow accessing the attributes of the runnable.

196 if hasattr(self.runnable, key):

197 return getattr(self.runnable, key)

@@ -198,11 +200,11 @@ class CUDAGraphWrapper:

200 f"cudagraph wrapper: {self.runnable}"

201 )

202

201 - def unwrap(self) -> Callable:

203 + def unwrap(self) -> Callable[..., Any]:

204 # in case we need to access the original runnable.

205 return self.runnable

206

205 - def __call__(self, *args, **kwargs):

207 + def __call__(self, *args: Any, **kwargs: Any) -> Any | None:

208 forward_context = get_forward_context()

209 batch_descriptor = forward_context.batch_descriptor

210 cudagraph_runtime_mode = forward_context.cudagraph_runtime_mode

vllm/compilation/decorators.py

+42 -20

@@ -6,8 +6,8 @@ import hashlib

6 import inspect

7 import os

8 import sys

9 -from collections.abc import Callable

10 -from typing import TypeVar, overload

9 +from collections.abc import Callable, Generator

10 +from typing import TYPE_CHECKING, Any, Literal, TypeVar, overload

11 from unittest.mock import patch

13 import torch

@@ -32,6 +32,14 @@ from vllm.utils.torch_utils import is_torch_equal_or_newer, supports_dynamo

33 from .monitor import start_monitoring_torch_compile

35 +if TYPE_CHECKING:

36 + # Only added on nightly/2.10 so wrap

37 + try:

38 + from torch._dynamo.package import SourceInfo

39 + except ImportError:

40 + # Fallback for old versions not supporting

41 + SourceInfo = Any

42 +

43 logger = init_logger(__name__)

45 IGNORE_COMPILE_KEY = "_ignore_compile_vllm"

@@ -59,7 +67,7 @@ def ignore_torch_compile(cls: _T) -> _T:

67 return cls

62 -def _should_ignore_torch_compile(cls) -> bool:

70 +def _should_ignore_torch_compile(cls: _T) -> bool:

71 """

72 Check if the class should be ignored for torch.compile.

73 """

@@ -224,7 +232,7 @@ def support_torch_compile(

232 return cls_decorator_helper

233

234

227 -def _model_hash_key(fn) -> str:

235 +def _model_hash_key(fn: Callable[..., Any]) -> str:

236 import vllm

237

238 sha256_hash = hashlib.sha256()

@@ -234,7 +242,9 @@ def _model_hash_key(fn) -> str:

242 return sha256_hash.hexdigest()

243

244

237 -def _verify_source_unchanged(source_info, vllm_config) -> None:

245 +def _verify_source_unchanged(

246 + source_info: "SourceInfo", vllm_config: VllmConfig

247 +) -> None:

248 from .caching import _compute_code_hash, _compute_code_hash_with_content

249

250 file_contents = {}

@@ -275,8 +285,12 @@ def _support_torch_compile(

285 setattr(cls, IGNORE_COMPILE_KEY, False)

286

287 def __init__(

278 - self, *, vllm_config: VllmConfig | None = None, prefix: str = "", **kwargs

279 - ):

288 + self: _T,

289 + *,

290 + vllm_config: VllmConfig | None = None,

291 + prefix: str = "",

292 + **kwargs: Any,

293 + ) -> None:

294 if vllm_config is None:

295 vllm_config = get_current_vllm_config()

296

@@ -309,13 +323,17 @@ def _support_torch_compile(

323

324 compilation_counter.num_models_seen += 1

325 self.compiled = False

312 - TorchCompileWithNoGuardsWrapper.__init__(self)

326 +

327 + # Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class

328 + TorchCompileWithNoGuardsWrapper.__init__(self) # type: ignore[arg-type]

329

330 cls.__init__ = __init__

331

316 - def _mark_dynamic_inputs(mod, type, *args, **kwargs):

317 - def mark_dynamic(arg, dims):

318 - if type == DynamicShapesType.UNBACKED:

332 + def _mark_dynamic_inputs(

333 + mod: _T, ds_type: DynamicShapesType, *args: Any, **kwargs: Any

334 + ) -> None:

335 + def mark_dynamic(arg: torch.Tensor, dims: list[int]) -> None:

336 + if ds_type == DynamicShapesType.UNBACKED:

337 if is_torch_equal_or_newer("2.10.0.dev"):

338 for dim in dims:

339 torch._dynamo.decorators.mark_unbacked(

@@ -326,7 +344,7 @@ def _support_torch_compile(

344 else:

345 torch._dynamo.mark_dynamic(arg, dims)

346

329 - sig = inspect.signature(mod.__class__.forward)

347 + sig = inspect.signature(mod.__class__.forward) # type: ignore[attr-defined]

348 bound_args = sig.bind(mod, *args, **kwargs)

349 bound_args.apply_defaults()

350 for k, dims in dynamic_arg_dims.items():

@@ -364,7 +382,7 @@ def _support_torch_compile(

382 else:

383 torch._dynamo.decorators.mark_unbacked(arg, dims)

384

367 - def __call__(self, *args, **kwargs):

385 + def __call__(self: _T, *args: Any, **kwargs: Any) -> Any:

386 # torch.compiler.is_compiling() means we are inside the compilation

387 # e.g. TPU has the compilation logic in model runner, so we don't

388 # need to compile the model inside.

@@ -444,7 +462,7 @@ def _support_torch_compile(

462 not envs.VLLM_USE_AOT_COMPILE

463 or self.vllm_config.compilation_config.backend == "eager"

464 )

447 - return TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs)

465 + return TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs) # type: ignore[arg-type]

466

467 # This is the path for the first compilation.

468 # the first compilation needs to have dynamic shapes marked

@@ -477,7 +495,7 @@ def _support_torch_compile(

495 # during Dynamo tracing, and their corresponding files

496 inline_call = InliningInstructionTranslator.inline_call_

497

480 - def patched_inline_call(self_):

498 + def patched_inline_call(self_: Any) -> Any:

499 code = self_.f_code

500 self.compilation_config.traced_files.add(code.co_filename)

501 return inline_call(self_)

@@ -535,7 +553,7 @@ def _support_torch_compile(

553 str(e),

554 )

555 else:

538 - output = TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs)

556 + output = TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs) # type: ignore[arg-type]

557

558 self.compiled = True

559 return output

@@ -545,7 +563,9 @@ def _support_torch_compile(

563

564

565 @contextlib.contextmanager

548 -def maybe_use_cudagraph_partition_wrapper(vllm_config: VllmConfig):

566 +def maybe_use_cudagraph_partition_wrapper(

567 + vllm_config: VllmConfig,

568 +) -> Generator[None, None, None]:

569 """

570 Context manager to set/unset customized cudagraph partition wrappers.

571

@@ -572,7 +592,9 @@ def maybe_use_cudagraph_partition_wrapper(vllm_config: VllmConfig):

592 current_platform.get_static_graph_wrapper_cls()

593 )

594

575 - def customized_cudagraph_wrapper(f, metadata: CUDAGraphWrapperMetadata):

595 + def customized_cudagraph_wrapper(

596 + f: Callable[..., Any], metadata: CUDAGraphWrapperMetadata

597 + ) -> Any:

598 partition_id = metadata.partition_index

599 num_partitions = metadata.num_partitions

600 return static_graph_wrapper_class(

@@ -600,7 +622,7 @@ def maybe_use_cudagraph_partition_wrapper(vllm_config: VllmConfig):

622

623

624 @contextlib.contextmanager

603 -def _torch27_patch_tensor_subclasses():

625 +def _torch27_patch_tensor_subclasses() -> Generator[None, None, None]:

626 """

627 Add support for using tensor subclasses (ie `BasevLLMParameter`, ect) when

628 using torch 2.7.0. This enables using weight_loader_v2 and the use of

@@ -614,7 +636,7 @@ def _torch27_patch_tensor_subclasses():

636 _ColumnvLLMParameter,

637 )

638

617 - def return_false(*args, **kwargs):

639 + def return_false(*args: Any, **kwargs: Any) -> Literal[False]:

640 return False

641

642 if version.parse("2.7") <= version.parse(torch.__version__) < version.parse("2.8"):

vllm/compilation/fix_functionalization.py

+5 -5

@@ -26,7 +26,7 @@ class FixFunctionalizationPass(VllmInductorPass):

26 """

28 @VllmInductorPass.time_and_log

29 - def __call__(self, graph: torch.fx.Graph):

29 + def __call__(self, graph: torch.fx.Graph) -> None:

30 # XPU does not support auto-functionalization yet.

31 # Will enable this when switch to vllm-xpu-kernels.

32 if current_platform.is_xpu():

@@ -179,7 +179,7 @@ class FixFunctionalizationPass(VllmInductorPass):

179 )

180 self.nodes_to_remove.clear()

181

182 - def _remove(self, node_or_nodes: torch.fx.Node | Iterable[torch.fx.Node]):

182 + def _remove(self, node_or_nodes: torch.fx.Node | Iterable[torch.fx.Node]) -> None:

183 """

184 Stage a node (or nodes) for removal at the end of the pass.

185 """

@@ -194,7 +194,7 @@ class FixFunctionalizationPass(VllmInductorPass):

194 node: torch.fx.Node,

195 mutated_args: dict[int, torch.fx.Node | str],

196 args: tuple[torch.fx.Node | str, ...] | None = None,

197 - ):

197 + ) -> None:

198 """

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

200 It also replaces the getitem users with the mutated arguments.

@@ -206,7 +206,7 @@ class FixFunctionalizationPass(VllmInductorPass):

206

207 def replace_users_with_mutated_args(

208 self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str]

209 - ):

209 + ) -> None:

210 """

211 Replace all getitem users of the auto-functionalized node with the

212 mutated arguments.

@@ -237,7 +237,7 @@ class FixFunctionalizationPass(VllmInductorPass):

237 graph: torch.fx.Graph,

238 node: torch.fx.Node,

239 args: tuple[torch.fx.Node | str, ...] | None = None,

240 - ):

240 + ) -> None:

241 """

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

243 If one of the kwargs is 'out', provide args directly,

vllm/compilation/inductor_pass.py

@@ -29,6 +29,9 @@ else:

29 Torch25CustomGraphPass as CustomGraphPass,

30 )

32 +# Re-export CustomGraphPass for external usage

33 +__all__ = ["CustomGraphPass"]

34 +

35 _pass_context = None

36 P = ParamSpec("P")

37 R = TypeVar("R")

vllm/compilation/noop_elimination.py

+2 -2

@@ -65,7 +65,7 @@ class NoOpEliminationPass(VllmInductorPass):

65 """

67 @VllmInductorPass.time_and_log

68 - def __call__(self, graph: torch.fx.Graph):

68 + def __call__(self, graph: torch.fx.Graph) -> None:

69 count = 0

70 # Remove no-op reshapes/views:

71 for node in graph.nodes:

@@ -117,7 +117,7 @@ class NoOpEliminationPass(VllmInductorPass):

117 2. The dimensions both correspond to the same SymInt

118 """

119 # Case 1

120 - return statically_known_true(dim == i_dim)

120 + return statically_known_true(dim == i_dim) # type: ignore[no-any-return]

121

122 def all_dims_equivalent(

123 self, dims: Iterable[int | SymInt], i_dims: Iterable[int | SymInt]

vllm/compilation/pass_manager.py

+19 -12

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

1 # SPDX-License-Identifier: Apache-2.0

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 import functools

4 +from collections.abc import Callable

5 +from typing import Any, ParamSpec, TypeVar

7 from torch import fx as fx

@@ -40,8 +42,11 @@ from .noop_elimination import NoOpEliminationPass

43 logger = init_logger(__name__)

45 +P = ParamSpec("P")

46 +R = TypeVar("R")

44 -def with_pattern_match_debug(fn):

48 +

49 +def with_pattern_match_debug(fn: Callable[P, R]) -> Callable[P, R]:

50 """

51 Function decorator that turns on inductor pattern match debug

52 for the duration of the call.

@@ -49,7 +54,7 @@ def with_pattern_match_debug(fn):

54 """

56 @functools.wraps(fn)

52 - def wrapper(*args, **kwargs):

57 + def wrapper(*args: P.args, **kwargs: P.kwargs) -> R:

58 if (debug_val := envs.VLLM_PATTERN_MATCH_DEBUG) is not None:

59 # optionally check rank here

60 with set_env_var("TORCHINDUCTOR_PATTERN_MATCH_DEBUG", debug_val):

@@ -59,7 +64,7 @@ def with_pattern_match_debug(fn):

64 return wrapper

62 -class PostGradPassManager(CustomGraphPass):

67 +class PostGradPassManager(CustomGraphPass): # type: ignore[misc]

68 """

69 The pass manager for post-grad passes.

70 It handles configuration, adding custom passes, and running passes.

@@ -74,11 +79,11 @@ class PostGradPassManager(CustomGraphPass):

79 This way, all passes operate on a functionalized graph.

80 """

77 - def __init__(self):

82 + def __init__(self) -> None:

83 self.passes: list[InductorPass] = []

85 @with_pattern_match_debug

81 - def __call__(self, graph: fx.Graph):

86 + def __call__(self, graph: fx.Graph) -> None:

87 VllmInductorPass.dump_prefix = 0 # reset dump index

89 compile_range = get_pass_context().compile_range

@@ -98,7 +103,7 @@ class PostGradPassManager(CustomGraphPass):

103 self.fix_functionalization(graph)

104 VllmInductorPass.dump_prefix = None # Cleanup index

105

101 - def configure(self, config: VllmConfig):

106 + def configure(self, config: VllmConfig) -> None:

107 self.pass_config = config.compilation_config.pass_config

108

109 # Set the current vllm config to allow tracing CustomOp instances

@@ -135,23 +140,25 @@ class PostGradPassManager(CustomGraphPass):

140 self.post_cleanup = PostCleanupPass(config)

141 self.fix_functionalization = FixFunctionalizationPass(config)

142

138 - def add(self, pass_: InductorPass):

143 + def add(self, pass_: InductorPass) -> None:

144 assert isinstance(pass_, InductorPass)

145 self.passes.append(pass_)

146

142 - def uuid(self):

147 + def uuid(self) -> str:

148 """

149 The PostGradPassManager is set as a custom pass in the Inductor and

150 affects compilation caching. Its uuid depends on the UUIDs of all

151 dependent passes and the pass config. See InductorPass for more info.

152 """

148 - state = {"pass_config": self.pass_config.compute_hash(), "passes": []}

153 + passes = []

154 +

155 + state: dict[str, Any] = {"pass_config": self.pass_config.compute_hash()}

156 for pass_ in self.passes:

150 - state["passes"].append(pass_.uuid())

151 - state["passes"].append(self.fix_functionalization.uuid())

157 + passes.append(pass_.uuid())

158 + passes.append(self.fix_functionalization.uuid())

159

160 # Include the compile range in the uuid to ensure that inductor

161 # recompiles the graph for the new dynamic compile range.

162 state["compile_range"] = str(get_pass_context().compile_range)

156 -

163 + state["passes"] = passes

164 return InductorPass.hash_dict(state)

vllm/compilation/piecewise_backend.py

+29 -15

@@ -86,27 +86,36 @@ class PiecewiseBackend:

86 self.to_be_compiled_ranges: set[Range] = set(self.compile_ranges)

88 # We only keep compilation management inside this class directly.

89 - for size in self.compile_sizes:

90 - range = Range(start=size, end=size)

91 - if range not in self.compile_ranges:

92 - self.range_entries[range] = RangeEntry(

93 - compile_range=range,

94 - )

95 - self.to_be_compiled_ranges.add(range)

89 + if self.compile_sizes is not None:

90 + for size in self.compile_sizes:

91 + if isinstance(size, str):

92 + assert size == "cudagraph_capture_sizes"

93 + raise NotImplementedError(

94 + "cudagraph_capture_sizes not supported in compile_sizes."

95 + "This should be handled in `post_init_cudagraph_sizes`."

96 + )

97 + else:

98 + assert isinstance(size, int)

99 + range = Range(start=size, end=size)

100 + if range not in self.compile_ranges:

101 + self.range_entries[range] = RangeEntry(

102 + compile_range=range,

103 + )

104 + self.to_be_compiled_ranges.add(range)

105

106 for range in self.compile_ranges:

107 self.range_entries[range] = RangeEntry(

108 compile_range=range,

109 )

110

102 - def check_for_ending_compilation(self):

111 + def check_for_ending_compilation(self) -> None:

112 if self.is_last_graph and not self.to_be_compiled_ranges:

113 # no specific sizes to compile

114 # save the hash of the inductor graph for the next run

115 self.vllm_backend.compiler_manager.save_to_file()

116 end_monitoring_torch_compile(self.vllm_config)

117

109 - def _fakify_args(self, args: list[Any]) -> list[Any]:

118 + def _fakify_args(self, args: tuple[Any, ...]) -> list[Any]:

119 # We need to pass fake example_inputs, otherwise torch.compile

120 # will fakify the example_inputs potentially causing some non dynamic

121 # dimension to be be duck shaped to other existing shapes that have hints

@@ -127,7 +136,9 @@ class PiecewiseBackend:

136 assert len(fake_example_inputs) == len(args)

137 return fake_example_inputs

138

130 - def _maybe_compile_for_range_entry(self, range_entry: RangeEntry, args) -> Any:

139 + def _maybe_compile_for_range_entry(

140 + self, range_entry: RangeEntry, args: tuple[Any, ...]

141 + ) -> Any:

142 if not range_entry.compiled:

143 range_entry.compiled = True

144 self.to_be_compiled_ranges.remove(range_entry.compile_range)

@@ -136,14 +147,14 @@ class PiecewiseBackend:

147 # fakify for range, real args for concrete size.

148 # For concrete size, we clear the shape env in

149 # compiler_manager.compile() so no need to fakify.

139 - args = (

150 + args_list = (

151 self._fakify_args(args)

152 if not range_entry.compile_range.is_single_size()

142 - else args

153 + else list(args)

154 )

155 range_entry.runnable = self.vllm_backend.compiler_manager.compile(

156 self.graph,

146 - args,

157 + args_list,

158 self.vllm_backend.inductor_config,

159 self.compilation_config,

160 compile_range=range_entry.compile_range,

@@ -153,10 +164,13 @@ class PiecewiseBackend:

164

165 self.check_for_ending_compilation()

166

156 - def _find_range_for_shape(self, runtime_shape: int) -> Range | None:

167 + def _find_range_for_shape(self, runtime_shape: int) -> RangeEntry | None:

168 # First we try to find the range entry for the concrete compile size

169 # If not found, we search for the range entry

170 # that contains the runtime shape.

171 + if self.compile_sizes is None:

172 + return None

173 +

174 if runtime_shape in self.compile_sizes:

175 return self.range_entries[Range(start=runtime_shape, end=runtime_shape)]

176 else:

@@ -165,7 +179,7 @@ class PiecewiseBackend:

179 return self.range_entries[range]

180 return None

181

168 - def __call__(self, *args) -> Any:

182 + def __call__(self, *args: Any) -> Any:

183 runtime_shape = args[self.sym_shape_indices[0]]

184 range_entry = self._find_range_for_shape(runtime_shape)

185

vllm/compilation/wrapper.py

+23 -13

@@ -4,9 +4,10 @@

4 import os

5 import sys

6 from abc import abstractmethod

7 +from collections.abc import Callable, Generator

8 from contextlib import contextmanager, nullcontext

9 from types import CodeType

9 -from typing import Any

10 +from typing import Any, ParamSpec, TypeVar

12 import torch

13 import torch._C._dynamo.guards

@@ -19,19 +20,26 @@ from vllm.utils.nvtx_pytorch_hooks import layerwise_nvtx_marker_context

21 logger = init_logger(__name__)

23 +R = TypeVar("R")

24 +P = ParamSpec("P")

23 -def _noop_add_global_state_guard(self, *args, **kwargs):

26 +

27 +def _noop_add_global_state_guard(

28 + self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any

29 +) -> None:

30 """No-op to skip the GLOBAL_STATE guard entirely"""

31 pass

28 -def _noop_add_torch_function_mode_stack_guard(self, *args, **kwargs):

34 +def _noop_add_torch_function_mode_stack_guard(

35 + self: torch._C._dynamo.guards.GuardManager, *args: Any, **kwargs: Any

36 +) -> None:

37 """No-op to skip the TORCH_FUNCTION_MODE_STACK guard entirely"""

38 pass

41 @contextmanager

34 -def _compilation_context():

42 +def _compilation_context() -> Generator[None, None, None]:

43 """Context manager for compilation settings and patches.

45 This manager:

@@ -88,13 +96,15 @@ class TorchCompileWithNoGuardsWrapper:

96 since we drop all guards.

97 """

91 - def check_invariants_and_forward(self, *args, **kwargs):

99 + def check_invariants_and_forward(self, *args: Any, **kwargs: Any) -> Any:

100 assert hasattr(self, "_check_shape_invariants")

101 self._check_shape_invariants(*args, **kwargs)

102

103 return self.forward(*args, **kwargs)

104

97 - def _call_with_optional_nvtx_range(self, callable_fn, *args, **kwargs):

105 + def _call_with_optional_nvtx_range(

106 + self, callable_fn: Callable[P, R], *args: P.args, **kwargs: P.kwargs

107 + ) -> Any:

108 if self.layerwise_nvtx_tracing_enabled:

109 args_list = list(args)

110 kwargs_dict = dict(kwargs)

@@ -108,7 +118,7 @@ class TorchCompileWithNoGuardsWrapper:

118 return ctx.result

119 return callable_fn(*args, **kwargs)

120

111 - def __init__(self):

121 + def __init__(self) -> None:

122 self.compiled = False

123

124 vllm_config = get_current_vllm_config()

@@ -192,9 +202,9 @@ class TorchCompileWithNoGuardsWrapper:

202

203 if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE:

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

195 - self._compiled_bytecode = None

205 + self._compiled_bytecode: CodeType | None = None

206

197 - def aot_compile(self, *args, **kwargs):

207 + def aot_compile(self, *args: Any, **kwargs: Any) -> Any:

208 if not hasattr(self._compiled_callable, "aot_compile"):

209 raise RuntimeError(

210 "aot_compile is not supported by the current configuration. "

@@ -203,7 +213,7 @@ class TorchCompileWithNoGuardsWrapper:

213 )

214 return self._compiled_callable.aot_compile((args, kwargs))

215

206 - def __call__(self, *args, **kwargs):

216 + def __call__(self, *args: Any, **kwargs: Any) -> Any:

217 if envs.VLLM_USE_BYTECODE_HOOK:

218 if (

219 self.vllm_config.compilation_config.mode

@@ -236,13 +246,13 @@ class TorchCompileWithNoGuardsWrapper:

246 )

247

248 @abstractmethod

239 - def forward(self, *args, **kwargs): ...

249 + def forward(self, *args: Any, **kwargs: Any) -> Any: ...

250

251 def original_code_object(self) -> CodeType:

252 """Return the original code object of the forward method."""

253 return self.__class__.forward.__code__

254

245 - def bytecode_hook(self, old_code: CodeType, new_code: CodeType):

255 + def bytecode_hook(self, old_code: CodeType, new_code: CodeType) -> None:

256 """Hook to save the compiled bytecode for direct execution."""

257 if old_code is not self.original_code_object():

258 return

@@ -299,7 +309,7 @@ class TorchCompileWithNoGuardsWrapper:

309 raise RuntimeError(msg)

310

311 @contextmanager

302 - def _dispatch_to_compiled_code(self):

312 + def _dispatch_to_compiled_code(self) -> Generator[None, None, None]:

313 # noqa: E501

314 """

315 Context manager to dispatch to internally compiled code for torch<2.8.

vllm/config/compilation.py

@@ -32,6 +32,9 @@ else:

33 logger = init_logger(__name__)

35 +# Explicitly exports Range

36 +__all__ = ["Range"]

37 +

39 class CompilationMode(enum.IntEnum):

40 """The compilation approach used for torch.compile-based compilation of the

vllm/distributed/device_communicators/pynccl_allocator.py

+1 -1

@@ -60,7 +60,7 @@ def is_symmetric_memory_tensor(tensor: torch.Tensor):

60 return False

63 -def set_graph_pool_id(graph_pool_id):

63 +def set_graph_pool_id(graph_pool_id: Any) -> None:

64 global _graph_pool_id

65 _graph_pool_id = graph_pool_id