[Misc][BE] Type coverage for vllm/compilation [3/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 #31744

Test Plan

mypy vllm/compilation

Test Result

Success: no issues found in 28 source files

---

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 coverage and API clarity across compilation passes with minimal logic adjustments.

• Add explicit return/param types, ParamSpec, and helper get_inputs methods across fusion passes (activation, RMSNorm, attention, collective, sequence parallelism, ROCm/AIter)
• Tighten function signatures (__init__, register, __call__, uuid, helpers like empty_*) and use precise tuple/list return types in pattern/replacement fns
• Safer device capability handling and minor no-op cleanup hooks; unify tracing wrappers (wrap_trace_fn), reshape conversions, and first-return-only helpers with typing
• Update decorators to import SourceInfo, mark dynamic/unbacked dims conditionally, and patch configs with typed contexts
• Minor typing fixes in distributed parallel_state, rotary embedding __all__ and cache key types

mypy: Success on 28 files (no issues).

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

---

Note

Improves type coverage and API clarity across vllm/compilation with minimal logic changes.

• Add explicit return/param types (incl. ParamSpec) and tighten signatures for __init__, register, __call__, uuid, and helper fns; pattern/replacement fns now return precise tuples
• Introduce typed get_inputs() helpers for patterns and unify tracing via wrap_trace_fn, view→reshape conversions, and no-op permute cleanup
• Safer FlashInfer handling (device capability may be None), workspace setup, and one-shot size checks; minor no-op cleanup hook in sequence parallelism
• Small typing fixes in distributed/parallel_state and rotary_embedding (__all__, cache key types)

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

[gemini-code-assist]

Code Review

This pull request focuses on improving type hint coverage across the vllm/compilation module, which is a valuable contribution to code quality and maintainability. The changes are extensive and well-executed, adding type hints to many functions and methods. I've identified one critical issue in vllm/compilation/collective_fusion.py where a replacement function in a pattern matcher returns a list of tensors instead of a single tensor, which will likely lead to a runtime error. Please address this issue.

[mergify]

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @Lucaskabela.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

[zou3519]

If someone asks: this PR does refactor get_inputs to be a methods on some of these classes but doesn't require it in the base class

[ProExpertProg]

Can we just set flashinfer_comm=None at the start and avoid the type: ignores?

[Lucaskabela]

If we set flashinfer_comm = None above this logic, we still need to add ignores for redef on L35 and L37 so not sure if it would save us much for code cleanliness

[ProExpertProg]

I thought that would just be reassignment not definition? Also I see the import is the same as the var, let's rename the import to _flashinfer_comm?

[hmellor]

Why do we need this?