[Misc][BE] Type coverage for vllm/compilation [2/3]

[Lucaskabela]

Purpose

We want to provide better type hint coverage in vllm/compilation to improve maintainability, readability, and reduce silent errors

This PR should be applied on top of #31554

Test Plan

mypy vllm/compilation

Test Result

---

Essential Elements of an Effective PR Description Checklist

• The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
• The test plan, such as providing test command.
• The test results, such as pasting the results comparison before and after, or e2e results
• (Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
• (Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

---

Note

Improves type safety and readability across compilation code with precise annotations and minor behavior clarifications.

• Broadly replaces untyped Callable/Any with concrete signatures (Callable[..., Any], Generator, ParamSpec, TypeVar, Literal) and adds return types throughout (backends, caching, cuda_graph, decorators, fix_functionalization, inductor_pass, noop_elimination, pass_manager, piecewise_backend, wrapper).
• piecewise_backend: validates compile_sizes (disallows "cudagraph_capture_sizes" here), supports None, returns RangeEntry from range lookup, and uses fakified inputs only for non-singleton ranges.
• decorators/wrapper: typed contexts and calls; preserve behavior while adding guard-related annotations and AOT/partition-wrapper hooks; minor mypy ignores where needed.
• pass_manager/inductor_pass: expose CustomGraphPass via __all__, typed manager methods, and include compile range in UUID state deterministically.
• caching: fully typed VllmSerializableFunction, safer graph pickling reducer, and typed co_name.
• cuda_graph: typed logging wrapper, wrapper methods, and callback signatures.
• config/compilation: export Range via __all__.
• Small logging/doc tweaks; no functional API changes expected aside from stricter validation and safer defaults.

^{Written by Cursor Bugbot for commit 58f1249. This will update automatically on new commits. Configure here.}

---

Note

^{Cursor Bugbot is generating a summary for commit 78d91f2. Configure here.}

---

Note

Improves type safety and clarity across the compilation stack with precise annotations and a few safe behavioral refinements.

• Broadly adds explicit types/returns to backends, caching, cuda_graph, decorators, fix_functionalization, inductor_pass, noop_elimination, pass_manager, piecewise_backend, and wrapper (with minimal mypy ignores)
• piecewise_backend: validates compile_sizes (disallows "cudagraph_capture_sizes" here), supports None, returns RangeEntry from range lookup, and only fakifies inputs for non-singleton ranges
• Deterministic/config-safe tweaks: re-export CustomGraphPass, include compile range in PostGradPassManager.uuid, export Range via config/compilation
• Safer serialization/logging and wrappers: typed GraphPickler reducer and co_name, typed CUDA graph wrapper callbacks, typed context managers and guard-handling in decorators/wrapper

^{Written by Cursor Bugbot for commit 78d91f2. This will update automatically on new commits. Configure here.}

[gemini-code-assist]

Code Review

This pull request continues the effort to improve type hint coverage in the vllm/compilation module. The changes are extensive and add type annotations to numerous functions, methods, and variables across multiple files. The added types are accurate and enhance code readability and maintainability. Several minor fixes are included, such as correcting return types and handling nullable arguments, which further improve the robustness of the code. The use of type: ignore is judicious and applied in contexts where type checkers might struggle with dynamic aspects of the code. Overall, this is a high-quality contribution that successfully improves type coverage.

[hmellor]

Why do we need to do this?

[hmellor]

Why do we need to do this?