vllm - 💡(How to fix) Fix [torch.compile] Proper PassManager infrastructure [1 participants]

Motivation

Currently, we bundle custom torch.compile passes using the PostGradPassManager. This manager is set as the current_platform.pass_key (default: "post_grad_custom_post_pass"), and individual passes within it are enabled/disabled using individual flags. This lacks scalability and control: every additional pass requires a new flag, which we often skip for utility passes (leaving no way to disable them). Platforms also have to choose a single pass_key as opposed to plugging into multiple points in the PyTorch compilation pipeline.

This inhibits debugging from the CLI, requiring manual code modifications. It also requires platforms to reimplement a lot of the pass management (enablement, cache hash, etc.). Another example is vLLM-Spyre, which currently plugs into "post_grad_custom_post_pass" for custom codegen transformations, colliding with the default vLLM pass manager.

Proposal

We should create a base PassManager class, which can handle invoking multiple passes and properly exposes a uuid() method to Inductor/AOTAutograd caching. We should also add a way to manually specify the full list of passes on the command line/in Python. Current flags can stick around and be used to produce this list by default. Passes should also be refactored to use 0-arg constructors by default (and read VllmConfig using get_current_vllm_config()) - that way OOT passes can integrate into this list more easily. This will also allow IR passes to move into the vllm/ir folder and remove their dependence on vLLM.

Timeline

This is not urgent and should probably wait until torch 2.12 when pre_grad_custom_pass becomes a proper CustomGraphPass. Depending on interest from OOT backends, this can be prioritized earlier.