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

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 9, 19:30 UTC

[Misc][BE] Turn on strict type coverage for vllm/compilation (#31756)

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

Lucas Kabela committed Jan 22, 2026 at 07:12 UTC 15e302dfced9e0fc31e4b6c2e8138b6957ed4740

11 files changed +121 -68

pyproject.toml

@@ -100,6 +100,13 @@ ignore_missing_imports = true

100 check_untyped_defs = true

101 follow_imports = "silent"

102

103 +[[tool.mypy.overrides]]

104 +module = "vllm.compilation.*"

105 +disallow_untyped_defs = true

106 +disallow_incomplete_defs = true

107 +warn_return_any = true

108 +follow_imports = "silent"

109 +

110 [tool.pytest.ini_options]

111 markers = [

112 "slow_test",

tests/compile/test_pass_manager.py

+1 -1

@@ -28,7 +28,7 @@ def test_bad_callable():

28 pass_manager.configure(config)

30 with pytest.raises(AssertionError):

31 - pass_manager.add(simple_callable)

31 + pass_manager.add(simple_callable) # type: ignore[arg-type]

34 # Pass that inherits from InductorPass

tools/pre_commit/mypy.py

+37 -3

@@ -77,6 +77,11 @@ EXCLUDE = [

77 "vllm/v1/attention/ops",

78 ]

80 +# Directories that should be checked with --strict

81 +STRICT_DIRS = [

82 + "vllm/compilation",

83 +]

84 +

86 def group_files(changed_files: list[str]) -> dict[str, list[str]]:

87 """

@@ -108,11 +113,17 @@ def group_files(changed_files: list[str]) -> dict[str, list[str]]:

113 return file_groups

114

115

116 +def is_strict_file(filepath: str) -> bool:

117 + """Check if a file should be checked with strict mode."""

118 + return any(filepath.startswith(strict_dir) for strict_dir in STRICT_DIRS)

119 +

120 +

121 def mypy(

122 targets: list[str],

123 python_version: str | None,

124 follow_imports: str | None,

125 file_group: str,

126 + strict: bool = False,

127 ) -> int:

128 """

129 Run mypy on the given targets.

@@ -124,6 +135,7 @@ def mypy(

135 follow_imports: Value for the --follow-imports option or None to use

136 the default mypy behavior.

137 file_group: The file group name for logging purposes.

138 + strict: If True, run mypy with --strict flag.

139

140 Returns:

141 The return code from mypy.

@@ -133,6 +145,8 @@ def mypy(

145 args += ["--python-version", python_version]

146 if follow_imports is not None:

147 args += ["--follow-imports", follow_imports]

148 + if strict:

149 + args += ["--strict"]

150 print(f"$ {' '.join(args)} {file_group}")

151 return subprocess.run(args + targets, check=False).returncode

152

@@ -149,9 +163,29 @@ def main():

163 for file_group, changed_files in file_groups.items():

164 follow_imports = None if ci and file_group == "" else "skip"

165 if changed_files:

152 - returncode |= mypy(

153 - changed_files, python_version, follow_imports, file_group

154 - )

166 + # Separate files into strict and non-strict groups

167 + strict_files = [f for f in changed_files if is_strict_file(f)]

168 + non_strict_files = [f for f in changed_files if not is_strict_file(f)]

169 +

170 + # Run mypy on non-strict files

171 + if non_strict_files:

172 + returncode |= mypy(

173 + non_strict_files,

174 + python_version,

175 + follow_imports,

176 + file_group,

177 + strict=False,

178 + )

179 +

180 + # Run mypy on strict files with --strict flag

181 + if strict_files:

182 + returncode |= mypy(

183 + strict_files,

184 + python_version,

185 + follow_imports,

186 + f"{file_group} (strict)",

187 + strict=True,

188 + )

189 return returncode

190

191

vllm/compilation/backends.py

+1 -1

@@ -68,7 +68,7 @@ def make_copy_and_call(

68 A wrapper function that copies inputs and calls the compiled function

69 """

71 - def copy_and_call(*args):

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

72 list_args = list(args)

73 for i, index in enumerate(sym_tensor_indices):

74 runtime_tensor = list_args[index]

vllm/compilation/caching.py

+12 -11

@@ -43,15 +43,15 @@ class StandaloneCompiledArtifacts:

43 split on attn)

44 """

46 - def __init__(self):

46 + def __init__(self) -> None:

47 # dict from submodule name to byte hash

48 - self.submodule_bytes = {}

48 + self.submodule_bytes: dict[str, str] = {}

49 # dict from byte hash to bytes

50 - self.submodule_bytes_store = {}

50 + self.submodule_bytes_store: dict[str, bytes] = {}

51 # dict from byte hash to loaded module

52 - self.loaded_submodule_store = {}

52 + self.loaded_submodule_store: dict[str, Any] = {}

54 - def insert(self, submod_name: str, shape: str, entry: bytes):

54 + def insert(self, submod_name: str, shape: str, entry: bytes) -> None:

55 hasher = hashlib.sha256()

56 hasher.update(entry)

57 hex_digest = hasher.hexdigest()

@@ -86,7 +86,7 @@ class StandaloneCompiledArtifacts:

86 self.submodule_bytes[f"{submod_name}_{shape}"]

87 ]

89 - def get_loaded(self, submod_name: str, shape: str):

89 + def get_loaded(self, submod_name: str, shape: str) -> Any:

90 logger.debug(

91 "getting artifact for submod %s with shape %s",

92 submod_name,

@@ -119,7 +119,7 @@ class StandaloneCompiledArtifacts:

119

120 from torch._inductor.standalone_compile import AOTCompiledArtifact

121

122 - def _load_entry(entry_bytes) -> AOTCompiledArtifact:

122 + def _load_entry(entry_bytes: bytes) -> AOTCompiledArtifact:

123 entry = pickle.loads(entry_bytes)

124 return AOTCompiledArtifact.deserialize(entry)

125

@@ -132,13 +132,13 @@ class StandaloneCompiledArtifacts:

132

133 logger.debug("loaded all %s submodules", self.num_artifacts())

134

135 - def __getstate__(self):

135 + def __getstate__(self) -> dict[str, dict[str, str] | dict[str, bytes]]:

136 return {

137 "submodule_bytes": self.submodule_bytes,

138 "submodule_bytes_store": self.submodule_bytes_store,

139 }

140

141 - def __setstate__(self, state):

141 + def __setstate__(self, state: dict[str, dict[str, Any]]) -> None:

142 self.submodule_bytes = state["submodule_bytes"]

143 self.submodule_bytes_store = state["submodule_bytes_store"]

144 self.loaded_submodule_store = {}

@@ -387,7 +387,7 @@ def reconstruct_serializable_fn_from_mega_artifact(

387 standalone_compile_artifacts.load_all()

388

389 submod_names = standalone_compile_artifacts.submodule_names()

390 - compiled_callables: dict[str, dict[str, Callable]] = {}

390 + compiled_callables: dict[str, dict[str, Callable[..., Any]]] = {}

391

392 for cache_key in standalone_compile_artifacts.submodule_bytes:

393 submod_name, shape_str = cache_key.rsplit("_", 1)

@@ -495,9 +495,10 @@ def _compute_code_hash_with_content(file_contents: dict[str, str]) -> str:

495 # e.g. exec(). We can't actually check these.

496 continue

497 hash_content.append(content)

498 - return safe_hash(

498 + result: str = safe_hash(

499 "\n".join(hash_content).encode(), usedforsecurity=False

500 ).hexdigest()

501 + return result

502

503

504 def _compute_code_hash(files: set[str]) -> str:

vllm/compilation/collective_fusion.py

+15 -13

@@ -30,19 +30,15 @@ from .vllm_inductor_pass import VllmInductorPass, VllmPatternMatcherPass

31 FP8_DTYPE = current_platform.fp8_dtype()

33 +flashinfer_comm: ModuleType | None = None

34 if find_spec("flashinfer"):

35 try:

35 - import flashinfer.comm as flashinfer_comm

36 + import flashinfer.comm as _flashinfer_comm

37 - flashinfer_comm: ModuleType | None = ( # type: ignore[no-redef]

38 - flashinfer_comm

39 - if hasattr(flashinfer_comm, "trtllm_allreduce_fusion")

40 - else None

41 - )

38 + if hasattr(_flashinfer_comm, "trtllm_allreduce_fusion"):

39 + flashinfer_comm = _flashinfer_comm

40 except ImportError:

43 - flashinfer_comm = None # type: ignore[assignment]

44 -else:

45 - flashinfer_comm = None # type: ignore[assignment]

41 + pass

43 logger = init_logger(__name__)

@@ -441,7 +437,7 @@ class AsyncTPPass(VllmPatternMatcherPass):

437 ):

438 return True

439 tp_size = get_tensor_model_parallel_world_size()

444 - return compile_range.is_single_size() and compile_range.end % tp_size == 0

440 + return bool(compile_range.is_single_size() and compile_range.end % tp_size == 0)

441

442 @VllmInductorPass.time_and_log

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

@@ -516,7 +512,7 @@ if flashinfer_comm is not None:

512 # Get one shot input size limit for the current world size

513 # for the current device capability

514 max_one_shot_size = _FI_ALLREDUCE_ONE_SHOT_MAX_SIZES_MB.get(

519 - device_capability, # type: ignore[arg-type]

515 + device_capability, # type: ignore[arg-type, unused-ignore]

516 {},

517 ).get(world_size, None)

518 # Use one shot if no max size is specified

@@ -666,6 +662,7 @@ class AllReduceRMSNormPattern(BasePattern):

662 ) -> tuple[torch.Tensor, torch.Tensor]:

663 residual = torch.zeros_like(input)

664 rms_result = torch.empty_like(input)

665 + assert flashinfer_comm is not None, "FlashInfer must be enabled"

666 allreduce = auto_functionalized(

667 flashinfer_trtllm_fused_allreduce_norm,

668 allreduce_in=input,

@@ -722,6 +719,7 @@ class AllReduceFusedAddRMSNormPattern(BasePattern):

719 def replacement(

720 residual: torch.Tensor, input: torch.Tensor, weight: torch.Tensor

721 ) -> tuple[torch.Tensor, torch.Tensor]:

722 + assert flashinfer_comm is not None, "FlashInfer must be enabled"

723 allreduce = auto_functionalized(

724 flashinfer_trtllm_fused_allreduce_norm,

725 allreduce_in=input,

@@ -800,6 +798,7 @@ class AllReduceFusedRMSNormStaticQuantFP8Pattern(BasePattern):

798 residual = torch.zeros_like(input)

799 result_rms = torch.empty_like(input)

800 result_quant = torch.empty_like(input, dtype=self.quant_dtype)

801 + assert flashinfer_comm is not None, "FlashInfer must be enabled"

802 allreduce = auto_functionalized(

803 flashinfer_trtllm_fused_allreduce_norm,

804 allreduce_in=input,

@@ -875,6 +874,7 @@ class AllReduceFusedAddRMSNormStaticQuantFP8Pattern(BasePattern):

874 scale: torch.Tensor,

875 ) -> tuple[torch.Tensor, torch.Tensor]:

876 result_quant = torch.empty_like(input, dtype=self.quant_dtype)

877 + assert flashinfer_comm is not None, "FlashInfer must be enabled"

878 allreduce = auto_functionalized(

879 flashinfer_trtllm_fused_allreduce_norm,

880 allreduce_in=input,

@@ -960,6 +960,7 @@ class AllReduceFusedRMSNormStaticQuantNVFP4Pattern(BasePattern):

960 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

961 residual = torch.zeros_like(input)

962 result_rms = torch.empty_like(input)

963 + assert flashinfer_comm is not None, "FlashInfer must be enabled"

964 allreduce = auto_functionalized(

965 flashinfer_trtllm_fused_allreduce_norm,

966 allreduce_in=input,

@@ -1055,6 +1056,7 @@ class AllReduceFusedAddRMSNormStaticQuantNVFP4Pattern(BasePattern):

1056 weight: torch.Tensor,

1057 input_global_scale: torch.Tensor,

1058 ) -> tuple[torch.Tensor, torch.Tensor, torch.Tensor]:

1059 + assert flashinfer_comm is not None, "FlashInfer must be enabled"

1060 allreduce = auto_functionalized(

1061 flashinfer_trtllm_fused_allreduce_norm,

1062 allreduce_in=input,

@@ -1131,7 +1133,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

1133 )

1134

1135 self.ipc_handles, workspace_tensor = (

1134 - flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion( # type: ignore[misc]

1136 + flashinfer_comm.trtllm_create_ipc_workspace_for_all_reduce_fusion(

1137 tp_rank=rank,

1138 tp_size=self.tp_size,

1139 max_token_num=self.max_token_num,

@@ -1204,7 +1206,7 @@ class AllReduceFusionPass(VllmPatternMatcherPass):

1206 if self.disabled:

1207 logger.warning_once("AllReduce fusion pass is disabled.")

1208 return False

1207 - return compile_range.end <= self.max_token_num

1209 + return bool(compile_range.end <= self.max_token_num)

1210

1211 @VllmInductorPass.time_and_log

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

vllm/compilation/compiler_interface.py

+6 -6

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

201

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

203 factors = get_inductor_factors()

204 - hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()[

205 - :10

206 - ]

204 + hash_str: str = safe_hash(

205 + str(factors).encode(), usedforsecurity=False

206 + ).hexdigest()[:10]

207 return hash_str

208

209 def initialize_cache(

@@ -319,9 +319,9 @@ class InductorAdaptor(CompilerInterface):

319

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

321 factors = get_inductor_factors()

322 - hash_str = safe_hash(str(factors).encode(), usedforsecurity=False).hexdigest()[

323 - :10

324 - ]

322 + hash_str: str = safe_hash(

323 + str(factors).encode(), usedforsecurity=False

324 + ).hexdigest()[:10]

325 return hash_str

326

327 def initialize_cache(

vllm/compilation/decorators.py

+17 -17

@@ -45,10 +45,10 @@ logger = init_logger(__name__)

46 IGNORE_COMPILE_KEY = "_ignore_compile_vllm"

48 -_T = TypeVar("_T", bound=type[nn.Module])

48 +_T = TypeVar("_T", bound=nn.Module)

51 -def ignore_torch_compile(cls: _T) -> _T:

51 +def ignore_torch_compile(cls: type[_T]) -> type[_T]:

52 """

53 A decorator to ignore support_torch_compile decorator

54 on the class. This is useful when a parent class has

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

68 return cls

71 -def _should_ignore_torch_compile(cls: _T) -> bool:

71 +def _should_ignore_torch_compile(cls: type[_T]) -> bool:

72 """

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

74 """

@@ -79,21 +79,21 @@ def _should_ignore_torch_compile(cls: _T) -> bool:

79 def support_torch_compile(

80 *,

81 enable_if: Callable[[VllmConfig], bool] | None = None,

82 -) -> Callable[[_T], _T]: ...

82 +) -> Callable[[type[_T]], type[_T]]: ...

85 @overload

86 def support_torch_compile(

87 *,

88 dynamic_arg_dims: dict[str, int | list[int]] | None,

89 -) -> Callable[[_T], _T]: ...

89 +) -> Callable[[type[_T]], type[_T]]: ...

92 @overload

93 def support_torch_compile(

94 *,

95 mark_unbacked_dims: dict[str, int | list[int]] | None,

96 -) -> Callable[[_T], _T]: ...

96 +) -> Callable[[type[_T]], type[_T]]: ...

99 @overload

@@ -101,21 +101,21 @@ def support_torch_compile(

101 *,

102 dynamic_arg_dims: dict[str, int | list[int]] | None,

103 mark_unbacked_dims: dict[str, int | list[int]] | None,

104 -) -> Callable[[_T], _T]: ...

104 +) -> Callable[[type[_T]], type[_T]]: ...

105

106

107 @overload

108 -def support_torch_compile(cls: _T) -> _T: ...

108 +def support_torch_compile(cls: type[_T]) -> type[_T]: ...

109

110

111 def support_torch_compile(

112 - cls: _T | None = None,

112 + cls: type[_T] | None = None,

113 *,

114 dynamic_arg_dims: dict[str, int | list[int]] | None = None,

115 mark_unbacked_dims: dict[str, int | list[int]] | None = None,

116 enable_if: Callable[[VllmConfig], bool] | None = None,

117 shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,

118 -) -> Callable[[_T], _T] | _T:

118 +) -> Callable[[type[_T]], type[_T]] | type[_T]:

119 """

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

121

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

182 errors.

183 """

184

185 - def cls_decorator_helper(cls: _T) -> _T:

185 + def cls_decorator_helper(cls: type[_T]) -> type[_T]:

186 # helper to pass `dynamic_arg_dims` to `_support_torch_compile`

187 # to avoid too much indentation for `_support_torch_compile`

188 if not hasattr(cls, "forward"):

@@ -263,12 +263,12 @@ def _verify_source_unchanged(

263

264

265 def _support_torch_compile(

266 - cls: _T,

266 + cls: type[_T],

267 dynamic_arg_dims: dict[str, int | list[int]],

268 mark_unbacked_dims: dict[str, int | list[int]] | None = None,

269 enable_if: Callable[[VllmConfig], bool] | None = None,

270 shape_invariants: Callable[..., None] = lambda *args, **kwargs: None,

271 -) -> _T:

271 +) -> type[_T]:

272 """

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

274 """

@@ -325,12 +325,12 @@ def _support_torch_compile(

325 self.compiled = False

326

327 # Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class

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

328 + TorchCompileWithNoGuardsWrapper.__init__(self)

329

330 cls.__init__ = __init__

331

332 def _mark_dynamic_inputs(

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

333 + mod: type[_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:

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

382 else:

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

384

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

385 + def __call__(self: type[_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.

@@ -564,7 +564,7 @@ def _support_torch_compile(

564 return output

565

566 # triggers VllmSerializableFunction.serialize()

567 - def save_aot_compiled_function(self):

567 + def save_aot_compiled_function(self: type[_T]) -> None:

568 if self.was_aot_compile_fn_loaded_from_disk:

569 logger.debug("AOT compiled function was loaded from cache, skipping save")

570 return

vllm/compilation/matcher_utils.py

+13 -9

@@ -141,15 +141,18 @@ class MatcherRotaryEmbedding(MatcherCustomOp):

141 key: torch.Tensor | None,

142 cos_sin_cache: torch.Tensor,

143 ) -> tuple[torch.Tensor, torch.Tensor | None]:

144 - return RotaryEmbedding.forward_static(

145 - positions,

146 - query,

147 - key,

148 - self.head_size,

149 - self.rotary_dim,

150 - cos_sin_cache,

151 - self.is_neox,

144 + result: tuple[torch.Tensor, torch.Tensor | None] = (

145 + RotaryEmbedding.forward_static(

146 + positions,

147 + query,

148 + key,

149 + self.head_size,

150 + self.rotary_dim,

151 + cos_sin_cache,

152 + self.is_neox,

153 + )

154 )

155 + return result

156

157

158 class MatcherRMSNorm(MatcherCustomOp):

@@ -275,9 +278,10 @@ class MatcherFusedAddRMSNorm(MatcherCustomOp):

278 weight: torch.Tensor,

279 residual: torch.Tensor,

280 ) -> tuple[torch.Tensor, torch.Tensor]:

278 - return RMSNorm.forward_static(

281 + result: tuple[torch.Tensor, torch.Tensor] = RMSNorm.forward_static(

282 input, self.epsilon, input.size(-1), self.model_dtype, weight, residual

283 )

284 + return result

285

286

287 class MatcherQuantFP8(MatcherCustomOp):

vllm/compilation/piecewise_backend.py

+8 -6

@@ -25,7 +25,7 @@ logger = init_logger(__name__)

25 class RangeEntry:

26 compile_range: Range

27 compiled: bool = False

28 - runnable: Callable = None # type: ignore

28 + runnable: Callable[..., Any] = None # type: ignore

31 class PiecewiseBackend:

@@ -38,7 +38,7 @@ class PiecewiseBackend:

38 sym_shape_indices: list[int],

39 vllm_backend: VllmBackend,

40 returns_tuple: bool,

41 - compiled_runnables: dict[str, Callable] | None = None,

41 + compiled_runnables: dict[str, Callable[..., Any]] | None = None,

42 ):

43 """

44 The backend for piecewise compilation.

@@ -138,8 +138,10 @@ class PiecewiseBackend:

138

139 self.on_compilation_complete = _on_compilation_complete_callback.get()

140

141 - def get_compiled_graph_wrapper(self, compiled_graph):

142 - def compiled_graph_wrapper(*args):

141 + def get_compiled_graph_wrapper(

142 + self, compiled_graph: Callable[..., Any]

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

144 + def compiled_graph_wrapper(*args: Any) -> Any:

145 graph_output = compiled_graph(*args)

146 # unpack the tuple if needed

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

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

165

166 def to_bytes(self) -> dict[str, bytes]:

167 class StandaloneCompiledArtifactsPickler(Pickler):

166 - def reducer_override(self, obj):

168 + def reducer_override(self, obj: object) -> Any:

169 if isinstance(obj, CachingAutotuner):

170 obj.prepare_for_pickle()

171 return pickle.loads, (

@@ -173,7 +175,7 @@ class PiecewiseBackend:

175 )

176 return NotImplemented

177

176 - def serialize(fn) -> bytes:

178 + def serialize(fn: Callable[..., Any]) -> bytes:

179 assert hasattr(fn, "serialize"), "fn must have serialize method"

180 with torch._functorch.config.patch("bundled_autograd_cache", True):

181 entry = fn.serialize()

vllm/compilation/sequence_parallelism.py

+4 -1

@@ -358,7 +358,10 @@ class SequenceParallelismPass(VllmPatternMatcherPass):

358 ):

359 return True

360 tp_size = get_tensor_model_parallel_world_size()

361 - return (compile_range.is_single_size()) and (compile_range.end % tp_size == 0)

361 + result: bool = (compile_range.is_single_size()) and (

362 + compile_range.end % tp_size == 0

363 + )

364 + return result

365

366 @VllmInductorPass.time_and_log

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