[Bugfix] [pytorch] Patch AOTAutogradCache._get_shape_env

[jamesjwu]

Description

After https://github.com/pytorch/pytorch/pull/151563/files, VLLM needs to patch an extra _get_shape_env() function when running inductor, as AOTAutogradCache now uses its own shape env function. The implementation is technically shared at torch._inductor.codecache.GuardedCache._get_shape_env, but adding it like this preserves backward compatibility with PyTorch 2.6.

Test Plan

The following program now runs when linking VLLM with pytorch main

from vllm import LLM, SamplingParams
import argparse

if __name__ == "__main__":
    parser = argparse.ArgumentParser(description='Script with a boolean flag')
    parser.add_argument('--model',
                        type=str,
                        default="facebook/opt-125m",
                        help='Model to use for generation')
    args = parser.parse_args()
    prompts = [
        "Hello, my name is",
        "The president of the United States is",
        "The capital of France is",
        "The future of AI is",
    ]
    sampling_params = SamplingParams(temperature=0.8, top_p=0.95)
    model = args.model

    compilation_config = {
        'level': 3,
        'compile_sizes': [1],
    }
    llm = LLM(model=model, compilation_config=compilation_config)
    outputs = llm.generate(prompts, sampling_params)

    for output in outputs:
        prompt = output.prompt
        generated_text = output.outputs[0].text
        print(f"Prompt: {prompt!r}, Generated text: {generated_text!r}")

[github-actions]

👋 Hi! Thank you for contributing to the vLLM project.

💬 Join our developer Slack at https://slack.vllm.ai to discuss your PR in #pr-reviews, coordinate on features in #feat- channels, or join special interest groups in #sig- channels.

Just a reminder: PRs would not trigger full CI run by default. Instead, it would only run fastcheck CI which starts running only a small and essential subset of CI tests to quickly catch errors. You can run other CI tests on top of those by going to your fastcheck build on Buildkite UI (linked in the PR checks section) and unblock them. If you do not have permission to unblock, ping simon-mo or khluu to add you in our Buildkite org.

Once the PR is approved and ready to go, your PR reviewer(s) can run CI to test the changes comprehensively before merging.

To run CI, PR reviewers can either: Add ready label to the PR or enable auto-merge.

🚀

[zou3519]

LGTM assuming the tests pass

[houseroad]

Can we add some unittest in the following PR to ensure we catch the issue in future?

[houseroad]

could you fix the pre-commit?

[jamesjwu]

• Fixed the pre-commit and gated the change to be backward compatible
• Will think about how to better unit test this as a followup: the thing is, the failure only happens on the latest version of torch (which VLLM does not pin), so in a way, all the existing unit tests would fail anyway once the pin updates. In fact, the unit test probably belongs on pytorch side, not here.