[BugFix] Patch inductor partitioning logic

[angelayi]

Purpose

Monkey-patches inductor 2.9 code to fix #26678

Test Plan

import os
from typing import Optional

import torch
import torch.nn as nn
from torch._dynamo.test_case import TestCase, run_tests
from torch._subclasses.fake_tensor import FakeTensorMode

from vllm import LLM, SamplingParams
from vllm.config import CompilationConfig, CompilationLevel, CUDAGraphMode
from tests.compile.backend import TestBackend
from vllm.config import CompilationConfig, PassConfig, VllmConfig, CompilationLevel


os.environ["TORCHINDUCTOR_FORCE_DISABLE_CACHES"] = "1"
os.environ["VLLM_DISABLE_COMPILE_CACHE"] = "1"
os.environ["VLLM_ENABLE_V1_MULTIPROCESSING"] = "1"
os.environ["VLLM_USE_V1"] = "1"
os.environ["VLLM_LOGGING_LEVEL"] = "DEBUG"
os.environ["VLLM_USE_STANDALONE_COMPILE"] = "1"

config = CompilationConfig(
    level=CompilationLevel.PIECEWISE,
    cudagraph_mode=CUDAGraphMode.FULL,
    # splitting_ops=[],
    custom_ops=['+quant_fp8'],
    use_inductor_graph_partition=True,
)

llm = LLM(
    model="RedHatAI/Meta-Llama-3.1-8B-Instruct-FP8",
    gpu_memory_utilization=0.6,
    max_model_len=3000,
    compilation_config=config,
    tensor_parallel_size=2,
    enforce_eager=False,
)

outputs = llm.generate(["Hello, my name is"], SamplingParams(temperature=0))

# Print the outputs.
print("-" * 50)
for output in outputs:
    prompt = output.prompt
    generated_text = output.outputs[0].text
    print(f"Prompt:    {prompt!r}")
    print(f"Output:    {generated_text!r}")
    print("-" * 60)

cc @zou3519 @ProExpertProg @BoyuanFeng

[gemini-code-assist]

Code Review

This pull request introduces a monkey-patch for PyTorch Inductor's partitioning logic to address a bug in torch version 2.9. The patch is applied conditionally based on the torch version. My review focuses on a misleading comment within the patch that contradicts the code's behavior, which could impact future maintainability. I've suggested a correction to make the comment accurate.

[gemini-code-assist]

The comment on line 31, so that we always return True, contradicts the function's implementation which can return False (as seen on line 96). This is misleading and could cause confusion for future maintenance. Please update the comment to accurately describe the patch's purpose, which appears to be reverting to a previous, correct behavior of should_partition.

Suggested change

	# Copied from torch._inductor.scheduler.Scheduler.should_partition. Patches
	# [this code](https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724)
	# so that we always return True.
	# This is a patched version of torch._inductor.scheduler.Scheduler.should_partition
	# that reverts to a prior implementation to fix a regression.
	# See: https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724

[ProExpertProg]

I think we should merge (some form of) #26116 first. We can make a new PR with 2.9, #26116, and this fix, and make sure tests in CI all pass (at least the compilation tests). After that we can merge #26116 and this PR in any order.

[zou3519]

I guess we should fix this as soon as we can in PyTorch and change this to just check torch==2.9.0, because someone can change GraphLowering._update_scheduler on PyTorch main :/. I'll remember to loop back to this post-PTC.

[BoyuanFeng]

+1 for patching only for torch==2.9.0

[zou3519]

lgtm, will let you two figure out the ordering of merging the PRs.

[BoyuanFeng]

Should we have a dedicated folder for pt2.9 monkey patch?

[ProExpertProg]

I think moving the import might just postpone the ci failure to the 2.9 PR but at least it'll resolve other tests there 👍

[BoyuanFeng]

would it work if we just use Scheduler.should_partition = should_partition_patched instead of _update_scheduler_patched?

[BoyuanFeng]

this does not hold..

[ProExpertProg]

We can fix in your PR, for main this is a no-op, in #26738 I manually made sure to use your approach