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

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, 15:15 UTC

[Misc][BE] Type coverage for vllm/compilation [1/3] (#31554)

Signed-off-by: Lucas Kabela <lucaskabela@meta.com>

Lucas Kabela committed Jan 6, 2026 at 17:37 UTC 873480d133f3e32743a0e187b03da1e67635bdc2

12 files changed +103 -85

vllm/compilation/backends.py

+13 -13

@@ -9,7 +9,7 @@ import operator

9 import os

10 import pprint

11 import time

12 -from collections.abc import Callable, Sequence

12 +from collections.abc import Callable, Generator, Sequence

13 from contextlib import contextmanager

14 from copy import deepcopy

15 from functools import partial

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

90 support int as key.

91 """

93 - def __init__(self, compilation_config: CompilationConfig):

93 + def __init__(self, compilation_config: CompilationConfig) -> None:

94 self.cache: dict[tuple[Range, int, str], Any] = dict()

95 self.is_cache_updated = False

96 self.compilation_config = compilation_config

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

100 return self.compiler.compute_hash(vllm_config)

101

102 @contextmanager

103 - def compile_context(self, compile_range: Range):

103 + def compile_context(self, compile_range: Range) -> Generator[None, None, None]:

104 """Provide compilation context for the duration of compilation to set

105 any torch global properties we want to scope to a single Inductor

106 compilation (e.g. partition rules, pass context)."""

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

115

116 def initialize_cache(

117 self, cache_dir: str, disable_cache: bool = False, prefix: str = ""

118 - ):

118 + ) -> None:

119 """

120 Initialize the cache directory for the compiler.

121

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

143 # do not use eval(), it is unsafe.

144 cache = ast.literal_eval(f.read())

145

146 - def check_type(value, ty):

146 + def check_type(value: Any, ty: type) -> None:

147 if not isinstance(value, ty):

148 raise TypeError(f"Expected {ty} but got {type(value)} for {value}")

149

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

165 cache_dir=cache_dir, disable_cache=disable_cache, prefix=prefix

166 )

167

168 - def save_to_file(self):

168 + def save_to_file(self) -> None:

169 if self.disable_cache or not self.is_cache_updated:

170 return

171 printer = pprint.PrettyPrinter(indent=4)

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

198 def compile(

199 self,

200 graph: fx.GraphModule,

201 - example_inputs,

201 + example_inputs: list[Any],

202 additional_inductor_config,

203 compilation_config: CompilationConfig,

204 compile_range: Range,

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

373 compile_submod_names: list[str],

374 vllm_config: VllmConfig,

375 vllm_backend: "VllmBackend",

376 - ):

376 + ) -> None:

377 super().__init__(module)

378 from torch._guards import detect_fake_mode

379

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

385 # When True, it annoyingly dumps the torch.fx.Graph on errors.

386 self.extra_traceback = False

387

388 - def run(self, *args):

388 + def run(self, *args: Any) -> Any:

389 # maybe instead just assert inputs are fake?

390 fake_args = [

391 self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t

@@ -467,7 +467,7 @@ model_is_encoder: bool = False

467

468

469 @contextmanager

470 -def set_model_tag(tag: str, is_encoder: bool = False):

470 +def set_model_tag(tag: str, is_encoder: bool = False) -> Generator[None, None, None]:

471 """Context manager to set the model tag."""

472 global model_tag

473 global model_is_encoder

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

521 vllm_config: VllmConfig,

522 prefix: str = "",

523 is_encoder: bool = False,

524 - ):

524 + ) -> None:

525 # if the model is initialized with a non-empty prefix,

526 # then usually it's enough to use that prefix,

527 # e.g. language_model, vision_model, etc.

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

558 # `torch.compile` is JIT compiled, so we don't need to

559 # do anything here

560

561 - def configure_post_pass(self):

561 + def configure_post_pass(self) -> None:

562 self.pass_manager.configure(self.vllm_config)

563

564 # Post-grad custom passes are run using the post_grad_custom_post_pass

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

580 self.inductor_config[self.pass_key] = self.pass_manager

581

582 def __call__(

583 - self, graph: fx.GraphModule, example_inputs

583 + self, graph: fx.GraphModule, example_inputs: Sequence[Any]

584 ) -> VllmSerializableFunction:

585 vllm_config = self.vllm_config

586 # Minimal hashing here with existing utilities, reused below.

vllm/compilation/collective_fusion.py

+7 -7

@@ -50,7 +50,7 @@ if hasattr(torch.ops._C, "scaled_fp4_quant"):

52 class BasePattern:

53 - def __init__(self, dtype: torch.dtype, device: str):

53 + def __init__(self, dtype: torch.dtype, device: str | None) -> None:

54 self.dtype = dtype

55 self.device = device

56 self.tp = get_tp_group()

@@ -637,7 +637,7 @@ class AllReduceRMSNormPattern(BasePattern):

637 self,

638 epsilon: float,

639 dtype: torch.dtype,

640 - device: str,

640 + device: str | None,

641 allreduce_params: FlashInferFusedAllReduceParams,

642 ):

643 super().__init__(dtype, device)

@@ -692,7 +692,7 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):

692 self,

693 epsilon: float,

694 dtype: torch.dtype,

695 - device: str,

695 + device: str | None,

696 allreduce_params: FlashInferFusedAllReduceParams,

697 ):

698 super().__init__(dtype, device)

@@ -759,7 +759,7 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):

759 self,

760 epsilon: float,

761 dtype: torch.dtype,

762 - device: str,

762 + device: str | None,

763 allreduce_params: FlashInferFusedAllReduceParams,

764 ):

765 super().__init__(dtype, device)

@@ -828,7 +828,7 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):

828 self,

829 epsilon: float,

830 dtype: torch.dtype,

831 - device: str,

831 + device: str | None,

832 allreduce_params: FlashInferFusedAllReduceParams,

833 ):

834 super().__init__(dtype, device)

@@ -902,7 +902,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):

902 self,

903 epsilon: float,

904 dtype: torch.dtype,

905 - device: str,

905 + device: str | None,

906 allreduce_params: FlashInferFusedAllReduceParams,

907 ):

908 super().__init__(dtype, device)

@@ -988,7 +988,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):

988 self,

989 epsilon: float,

990 dtype: torch.dtype,

991 - device: str,

991 + device: str | None,

992 allreduce_params: FlashInferFusedAllReduceParams,

993 ):

994 super().__init__(dtype, device)

vllm/compilation/compiler_interface.py

+24 -24

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

32 def initialize_cache(

33 self, cache_dir: str, disable_cache: bool = False, prefix: str = ""

34 - ):

34 + ) -> None:

35 """

36 when the vLLM process uses `cache_dir` as the cache directory,

37 the compiler should initialize itself with the cache directory,

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

66 compiler_config: dict[str, Any],

67 compile_range: Range,

68 key: str | None = None,

69 - ) -> tuple[Callable | None, Any | None]:

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

70 """

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

72 with a range. The `compile_range` specifies the range of the inputs,

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

100 example_inputs: list[Any],

101 graph_index: int,

102 compile_range: Range,

103 - ) -> Callable:

103 + ) -> Callable[..., Any]:

104 """

105 Load the compiled function from the handle.

106 Raises an error if the handle is invalid.

@@ -138,13 +138,13 @@ class AlwaysHitShapeEnv:

138 def __init__(self) -> None:

139 self.guards: list[Any] = []

140

141 - def evaluate_guards_expression(self, *args, **kwargs):

141 + def evaluate_guards_expression(self, *args: Any, **kwargs: Any) -> Literal[True]:

142 return True

143

144 - def get_pruned_guards(self, *args, **kwargs):

144 + def get_pruned_guards(self, *args: Any, **kwargs: Any) -> list[Any]:

145 return []

146

147 - def produce_guards_expression(self, *args, **kwargs):

147 + def produce_guards_expression(self, *args: Any, **kwargs: Any) -> Literal[""]:

148 return ""

149

150

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

193

194 name = "inductor_standalone"

195

196 - def __init__(self, save_format: Literal["binary", "unpacked"]):

196 + def __init__(self, save_format: Literal["binary", "unpacked"]) -> None:

197 self.save_format = save_format

198

199 def compute_hash(self, vllm_config: VllmConfig) -> str:

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

205

206 def initialize_cache(

207 self, cache_dir: str, disable_cache: bool = False, prefix: str = ""

208 - ):

208 + ) -> None:

209 self.cache_dir = cache_dir

210

211 def compile(

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

215 compiler_config: dict[str, Any],

216 compile_range: Range,

217 key: str | None = None,

218 - ) -> tuple[Callable | None, Any | None]:

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

219 compilation_counter.num_inductor_compiles += 1

220 current_config = {}

221 if compiler_config is not None:

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

252 example_inputs: list[Any],

253 graph_index: int,

254 compile_range: Range,

255 - ) -> Callable:

255 + ) -> Callable[..., Any]:

256 assert isinstance(handle, tuple)

257 assert isinstance(handle[0], str)

258 assert isinstance(handle[1], str)

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

264

265 returns_tuple = graph_returns_tuple(graph)

266

267 - def compiled_graph_wrapper(*args):

267 + def compiled_graph_wrapper(*args: Any) -> tuple[Any, ...] | Any:

268 graph_output = inductor_compiled_graph(*args)

269 # unpack the tuple if needed

270 # TODO(rzou): the implication is that we're not

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

293

294 def initialize_cache(

295 self, cache_dir: str, disable_cache: bool = False, prefix: str = ""

296 - ):

296 + ) -> None:

297 self.cache_dir = cache_dir

298 self.prefix = prefix

299 self.base_cache_dir = cache_dir[: -len(prefix)] if prefix else cache_dir

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

317 compiler_config: dict[str, Any],

318 compile_range: Range,

319 key: str | None = None,

320 - ) -> tuple[Callable | None, Any | None]:

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

321 compilation_counter.num_inductor_compiles += 1

322 from torch._inductor.compile_fx import compile_fx

323

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

348 original_load = FxGraphCache.load

349 original_load_name = "torch._inductor.codecache.FxGraphCache.load"

350

351 - def hijack_load(*args, **kwargs):

351 + def hijack_load(*args: Any, **kwargs: Any) -> Any:

352 inductor_compiled_graph = original_load(*args, **kwargs)

353 nonlocal file_path

354 compiled_fn = inductor_compiled_graph.current_callable

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

375 # function renamed in 2.6

376 original_load_name = None

377

378 - def hijacked_compile_fx_inner(*args, **kwargs):

378 + def hijacked_compile_fx_inner(*args: Any, **kwargs: Any) -> Any:

379 output = torch._inductor.compile_fx.compile_fx_inner(*args, **kwargs)

380 nonlocal hash_str

381 inductor_compiled_graph = output

@@ -401,13 +401,13 @@ class InductorAdaptor(CompilerInterface):

401 hash_str = inductor_compiled_graph._fx_graph_cache_key

402 return output

403

404 - def hijack_compiled_fx_graph_hash(*args, **kwargs):

404 + def hijack_compiled_fx_graph_hash(*args: Any, **kwargs: Any) -> Any:

405 out = compiled_fx_graph_hash(*args, **kwargs)

406 nonlocal hash_str

407 hash_str = out[0]

408 return out

409

410 - def _check_can_cache(*args, **kwargs):

410 + def _check_can_cache(*args: Any, **kwargs: Any) -> None:

411 # no error means it can be cached.

412 # Inductor refuses to cache the graph outside of Dynamo

413 # tracing context, and also disables caching for graphs

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

513 example_inputs: list[Any],

514 graph_index: int,

515 compile_range: Range,

516 - ) -> Callable:

516 + ) -> Callable[..., Any]:

517 assert isinstance(handle, tuple)

518 assert isinstance(handle[0], str)

519 assert isinstance(handle[1], str)

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

572 returns_tuple = graph_returns_tuple(graph)

573

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

575 - def compiled_graph(*args):

575 + def compiled_graph(*args: Any) -> tuple[Any, ...] | Any:

576 # convert args to list

577 list_args = list(args)

578 graph_output = inductor_compiled_graph(list_args)

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

584

585 return compiled_graph

586

587 - def metrics_context(self) -> contextlib.AbstractContextManager:

587 + def metrics_context(self) -> contextlib.AbstractContextManager[Any]:

588 """

589 This method returns the Dynamo metrics context (if it exists,

590 otherwise a null context). It is used by various compile components.

@@ -603,12 +603,12 @@ class InductorAdaptor(CompilerInterface):

603 if is_torch_equal_or_newer("2.6"):

604 import torch._dynamo.utils

605

606 - return torch._dynamo.utils.get_metrics_context()

606 + return torch._dynamo.utils.get_metrics_context() # type: ignore[no-any-return]

607 else:

608 return contextlib.nullcontext()

609

610

611 -def set_inductor_config(config, compile_range: Range):

611 +def set_inductor_config(config: dict[str, Any], compile_range: Range) -> None:

612 if compile_range.is_single_size():

613 # for a specific batch size, tuning triton kernel parameters

614 # can be beneficial

@@ -618,7 +618,7 @@ def set_inductor_config(config, compile_range: Range):

618 )

619

620

621 -def set_functorch_config():

621 +def set_functorch_config() -> None:

622 torch._functorch.config.bundled_autograd_cache = False

623

624

@@ -632,7 +632,7 @@ class EagerAdaptor(CompilerInterface):

632 compiler_config: dict[str, Any],

633 compile_range: Range,

634 key: str | None = None,

635 - ) -> tuple[Callable | None, Any | None]:

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

636 compilation_counter.num_eager_compiles += 1

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

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

vllm/compilation/counter.py

+3 -1

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

4 import copy

5 import dataclasses

6 +from collections.abc import Generator

7 from contextlib import contextmanager

8 +from typing import Any

11 @dataclasses.dataclass

@@ -34,7 +36,7 @@ class CompilationCounter:

36 return copy.deepcopy(self)

38 @contextmanager

37 - def expect(self, **kwargs):

39 + def expect(self, **kwargs: Any) -> Generator[None, None, None]:

40 old = self.clone()

41 yield

42 for k, v in kwargs.items():

vllm/compilation/cuda_graph.py

@@ -219,6 +219,7 @@ class CUDAGraphWrapper:

219 # runtime modes.

220 return self.runnable(*args, **kwargs)

221

222 + assert batch_descriptor is not None

223 if batch_descriptor not in self.concrete_cudagraph_entries:

224 # create a new entry for this batch descriptor

225 self.concrete_cudagraph_entries[batch_descriptor] = CUDAGraphEntry(

vllm/compilation/fx_utils.py

+3 -2

@@ -7,10 +7,11 @@ from collections.abc import Iterable, Iterator

7 from torch import fx

8 from torch._higher_order_ops.auto_functionalize import auto_functionalized

9 from torch._ops import OpOverload, OpOverloadPacket

10 +from torch.fx.node import Target

12 -def is_func(node: fx.Node, target) -> bool:

13 - return node.op == "call_function" and node.target == target

13 +def is_func(node: fx.Node, target: Target) -> bool:

14 + return bool(node.op == "call_function" and node.target == target)

17 def is_auto_func(node: fx.Node, op: OpOverload) -> bool:

vllm/compilation/inductor_pass.py

+15 -11

@@ -8,9 +8,9 @@ import hashlib

8 import inspect

9 import json

10 import types

11 -from collections.abc import Callable

11 +from collections.abc import Callable, Generator

12 from contextlib import contextmanager

13 -from typing import TYPE_CHECKING, Any

13 +from typing import TYPE_CHECKING, Any, ParamSpec, TypeVar

15 import torch

16 from torch import fx

@@ -30,6 +30,8 @@ else:

30 )

32 _pass_context = None

33 +P = ParamSpec("P")

34 +R = TypeVar("R")

37 class PassContext:

@@ -44,7 +46,7 @@ def get_pass_context() -> PassContext:

48 @contextmanager

47 -def pass_context(compile_range: Range):

49 +def pass_context(compile_range: Range) -> Generator[None, None, None]:

50 """A context manager that stores the current pass context,

51 usually it is a list of sizes to specialize.

52 """

@@ -57,7 +59,7 @@ def pass_context(compile_range: Range):

59 _pass_context = prev_context

60 -class InductorPass(CustomGraphPass):

62 +class InductorPass(CustomGraphPass): # type: ignore[misc]

63 """

64 A custom graph pass that uses a hash of its source as the UUID.

65 This is defined as a convenience and should work in most cases.

@@ -73,7 +75,7 @@ class InductorPass(CustomGraphPass):

75 return InductorPass.hash_source(self)

77 @staticmethod

76 - def hash_source(*srcs: str | Any):

78 + def hash_source(*srcs: str | Any) -> str:

79 """

80 Utility method to hash the sources of functions or objects.

81 :param srcs: strings or objects to add to the hash.

@@ -93,7 +95,7 @@ class InductorPass(CustomGraphPass):

95 return hasher.hexdigest()

97 @staticmethod

96 - def hash_dict(dict_: dict[Any, Any]):

98 + def hash_dict(dict_: dict[Any, Any]) -> str:

99 """

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

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

@@ -101,7 +103,7 @@ class InductorPass(CustomGraphPass):

103 encoded = json.dumps(dict_, sort_keys=True).encode("utf-8")

104 return hashlib.sha256(encoded).hexdigest()

105

104 - def is_applicable_for_range(self, compile_range: Range):

106 + def is_applicable_for_range(self, compile_range: Range) -> bool:

107 return True

108

109

@@ -111,25 +113,27 @@ class CallableInductorPass(InductorPass):

113 implementation of the UUID.

114 """

115

114 - def __init__(self, callable: Callable[[fx.Graph], None], uuid: Any | None = None):

116 + def __init__(

117 + self, callable: Callable[[fx.Graph], None], uuid: Any | None = None

118 + ) -> None:

119 self.callable = callable

120 self._uuid = self.hash_source(callable) if uuid is None else uuid

121

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

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

123 self.callable(graph)

124

125 def uuid(self) -> Any:

126 return self._uuid

127

128

125 -def enable_fake_mode(fn: Callable[..., Any]) -> Callable[..., Any]:

129 +def enable_fake_mode(fn: Callable[P, R]) -> Callable[P, R]:

130 """

131 Applies a FakeTensorMode context. This is useful when you don't want to

132 create or run things with real tensors.

133 """

134

135 @functools.wraps(fn)

132 - def fn_new(*args, **kwargs) -> Any:

136 + def fn_new(*args: P.args, **kwargs: P.kwargs) -> R:

137 with torch._guards.tracing(None), unset_fake_temporarily(), FakeTensorMode():

138 result = fn(*args, **kwargs)

139

vllm/compilation/monitor.py

+4 -4

@@ -12,7 +12,7 @@ context_manager = None

12 torch_compile_start_time: float = 0.0

15 -def start_monitoring_torch_compile(vllm_config: VllmConfig):

15 +def start_monitoring_torch_compile(vllm_config: VllmConfig) -> None:

16 global torch_compile_start_time

17 torch_compile_start_time = time.time()

@@ -28,7 +28,7 @@ def start_monitoring_torch_compile(vllm_config: VllmConfig):

28 context_manager.__enter__()

31 -def end_monitoring_torch_compile(vllm_config: VllmConfig):

31 +def end_monitoring_torch_compile(vllm_config: VllmConfig) -> None:

32 compilation_config: CompilationConfig = vllm_config.compilation_config

33 if compilation_config.mode == CompilationMode.VLLM_COMPILE:

34 logger.info_once(

@@ -45,7 +45,7 @@ def end_monitoring_torch_compile(vllm_config: VllmConfig):

45 cudagraph_capturing_enabled: bool = True

48 -def validate_cudagraph_capturing_enabled():

48 +def validate_cudagraph_capturing_enabled() -> None:

49 # used to monitor whether a cudagraph capturing is legal at runtime.

50 # should be called before any cudagraph capturing.

51 # if an illegal cudagraph capturing happens, raise an error.

@@ -57,6 +57,6 @@ def validate_cudagraph_capturing_enabled():

57 )

60 -def set_cudagraph_capturing_enabled(enabled: bool):

60 +def set_cudagraph_capturing_enabled(enabled: bool) -> None:

61 global cudagraph_capturing_enabled

62 cudagraph_capturing_enabled = enabled

vllm/compilation/partition_rules.py

+4 -1

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

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

4 import contextlib

5 +from collections.abc import Generator

7 import torch

@@ -38,7 +39,9 @@ def should_split(node: torch.fx.Node, splitting_ops: list[str]) -> bool:

41 @contextlib.contextmanager

41 -def inductor_partition_rule_context(splitting_ops: list[str]):

42 +def inductor_partition_rule_context(

43 + splitting_ops: list[str] | None,

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

45 """Context manager to temporarily register Inductor partition rules.

47 Registers custom partition rules for specified operators, forcing the

vllm/compilation/sequence_parallelism.py

+7 -7

@@ -41,8 +41,8 @@ class _SequenceParallelPatternHelper:

41 self,

42 epsilon: float,

43 dtype: torch.dtype,

44 - device: str,

45 - ):

44 + device: str | None,

45 + ) -> None:

46 self.epsilon = epsilon

47 self.dtype = dtype

48 self.device = device

@@ -64,7 +64,7 @@ class _SequenceParallelPatternHelper:

66 class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):

67 - def __init__(self, epsilon: float, dtype: torch.dtype, device: str):

67 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None) -> None:

68 super().__init__(epsilon, dtype, device)

69 self.rmsnorm_matcher = MatcherRMSNorm(epsilon)

@@ -74,7 +74,7 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):

75 return [input, arg3_1]

77 - def register(self, pm_pass: PatternMatcherPass):

77 + def register(self, pm_pass: PatternMatcherPass) -> None:

78 def pattern(

79 input: torch.Tensor,

80 arg3_1: torch.Tensor,

@@ -100,7 +100,7 @@ class FirstAllReduceRMSNormPattern(_SequenceParallelPatternHelper):

100

101

102 class MiddleAllReduceRMSNormPattern(_SequenceParallelPatternHelper):

103 - def __init__(self, epsilon: float, dtype: torch.dtype, device: str):

103 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None):

104 super().__init__(epsilon, dtype, device)

105 self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)

106

@@ -162,7 +162,7 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):

162 self,

163 epsilon: float,

164 dtype: torch.dtype,

165 - device: str,

165 + device: str | None,

166 ):

167 super().__init__(epsilon, dtype, device)

168 self.rmsnorm_matcher = MatcherRMSNorm(epsilon)

@@ -203,7 +203,7 @@ class FirstAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):

203

204

205 class MiddleAllReduceRMSNormStaticFP8Pattern(_SequenceParallelPatternHelper):

206 - def __init__(self, epsilon: float, dtype: torch.dtype, device: str):

206 + def __init__(self, epsilon: float, dtype: torch.dtype, device: str | None):

207 super().__init__(epsilon, dtype, device)

208 self.rmsnorm_matcher = MatcherFusedAddRMSNorm(epsilon)

209 self.quant_matcher = MatcherQuantFP8(kFp8StaticTensorSym)

vllm/compilation/torch25_custom_graph_pass.py

+3 -3

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

1 # SPDX-License-Identifier: Apache-2.0

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

3 from abc import ABC, abstractmethod

4 -from typing import Any

4 +from typing import Any, NoReturn

6 import torch

@@ -29,14 +29,14 @@ class Torch25CustomGraphPass(ABC): # noqa (redefinition)

29 Return None to skip inductor code caching entirely.

30 """

32 - def __getstate__(self):

32 + def __getstate__(self) -> Any | None:

33 """

34 Pickling is used instead of uuid() in torch<2.6. Just return uuid()

35 to enable subclasses to only have to implement uuid.

36 """

37 return self.uuid()

39 - def __setstate__(self, state):

39 + def __setstate__(self, state: Any) -> NoReturn:

40 raise ValueError(

41 "Cannot unpickle CustomGraphPass because pickling"

42 " is used for cache key uuid. Use torch>=2.6 with"

vllm/compilation/vllm_inductor_pass.py

+19 -12

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

3 import functools

4 import operator

5 import time

6 +from collections.abc import Callable

7 from dataclasses import dataclass

8 from typing import ClassVar

@@ -43,13 +44,17 @@ class VllmInductorPass(InductorPass):

44 )

45 self.pass_config = config.compilation_config.pass_config

46 self.model_dtype = config.model_config.dtype if config.model_config else None

46 - self.device = config.device_config.device if config.device_config else None

47 + self.device: str | None = (

48 + config.device_config.device if config.device_config else None

49 + )

50 self.pass_name = self.__class__.__name__

52 @staticmethod

50 - def time_and_log(call_fn):

53 + def time_and_log(

54 + call_fn: Callable[["VllmInductorPass", torch.fx.Graph], None],

55 + ) -> Callable[["VllmInductorPass", torch.fx.Graph], None]:

56 @functools.wraps(call_fn)

52 - def wrapped(self: VllmInductorPass, graph: torch.fx.Graph):

57 + def wrapped(self: VllmInductorPass, graph: torch.fx.Graph) -> None:

58 self.begin()

59 self.dump_graph(graph, "before")

60 call_fn(self, graph)

@@ -58,17 +63,17 @@ class VllmInductorPass(InductorPass):

64 return wrapped

61 - def dump_graph(self, graph: torch.fx.Graph, stage: str):

66 + def dump_graph(self, graph: torch.fx.Graph, stage: str) -> None:

67 i = VllmInductorPass.dump_prefix

68 i_str = "" if i is None else f".{i}"

69 lazy_format_graph_code(

70 f"post_grad{i_str}.{self.pass_name}.{stage}", graph.owning_module

71 )

68 - def begin(self):

73 + def begin(self) -> None:

74 self._start_time = time.perf_counter_ns()

71 - def end_and_log(self):

76 + def end_and_log(self) -> None:

77 self._end_time = time.perf_counter_ns()

78 duration_ms = float(self._end_time - self._start_time) / 1.0e6

79 logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms)

@@ -92,12 +97,14 @@ class VllmPatternMatcherPass(VllmInductorPass):

98 def _replace_op_overloads(self, string: str) -> str:

99 """Replace <OpOverload(..., ...)> with nicer formulations"""

95 - return self._OP_OVERLOAD_PATTERN.sub(

96 - lambda m: f"torch.ops.{m.group(1)}.{m.group(2)}",

97 - string,

100 + return str(

101 + self._OP_OVERLOAD_PATTERN.sub(

102 + lambda m: f"torch.ops.{m.group(1)}.{m.group(2)}",

103 + string,

104 + )

105 )

106

100 - def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass):

107 + def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None:

108 """

109 If debug dumping is enabled, dump the Inductor pattern-matcher patterns

110 into the debug_dump_path folder next to the dumped fx graphs.

@@ -165,9 +172,9 @@ class VllmPatternMatcherPass(VllmInductorPass):

172

173

174 class PrinterInductorPass(VllmInductorPass):

168 - def __init__(self, name: str, config: VllmConfig):

175 + def __init__(self, name: str, config: VllmConfig) -> None:

176 super().__init__(config)

177 self.name = name

178

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

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

180 self.dump_graph(graph, self.name)