A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 14, 05:30 UTC
[BugFix] Patch inductor partitioning logic (#26735)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Oct 14, 2025 at 22:04 UTC
7cfa420f4927a9e3fa4f533b2169d51597329a96
1 file changed
+118
vllm/env_override.py
+118
@@ -3,6 +3,7 @@
3
import os
4
5
import torch
6
+from packaging import version
7
8
from vllm.logger import init_logger
9
@@ -21,3 +22,120 @@ os.environ["PYTORCH_NVML_BASED_CUDA_CHECK"] = "1"
22
os.environ["TORCHINDUCTOR_COMPILE_THREADS"] = "1"
23
# see https://github.com/vllm-project/vllm/issues/10619
24
torch._inductor.config.compile_threads = 1
25
+
26
+
27
+# ========================================
28
+# torch 2.9 Inductor Scheduler monkeypatch
29
+# ========================================
30
+# This change monkeypatches a function in Inductor to work around the following
31
+# bug: https://github.com/vllm-project/vllm/issues/26678
32
+#
33
+# The bug occurs when `use_inductor_graph_partition` is turned on and there
34
+# exists operators inside of `splitting_ops` that have an in-place mutation. In
35
+# vllm, this specifically occurs on the operator
36
+# vllm.unified_attention_with_output. In this case, inductor does not populate
37
+# the inductor IR's `origin_node` field, causing an assertion error when trying
38
+# to access the node's `origin_node` field.
39
+#
40
+# So, we will monkeypatch torch._inductor.scheduler.Scheduler.should_partition
41
+# so that it does not access the inductor IR node's `origin_node` field and just
42
+# returns True if a node is registered as having a custom partition function.
43
+# This is ok for now since vllm's implementation of the custom partition
44
+# functions just return True.
45
+# ========================================
46
+
47
+
48
+def should_partition_patched(self, node, should_log: bool = False) -> bool:
49
+ # This is a patched version of
50
+ # torch._inductor.scheduler.Scheduler.should_partition that modifies
51
+ # the following piece of code so that we always return True:
52
+ # https://github.com/pytorch/pytorch/blob/ecb53078faf86ca1b33277df33b82985675bb011/torch/_inductor/scheduler.py#L4712-L4724
53
+ """Return True if we should partition the inductor graph on this node"""
54
+
55
+ import torch._inductor.ir as ir
56
+ from torch._inductor.scheduler import (
57
+ BaseSchedulerNode,
58
+ FusedSchedulerNode,
59
+ _custom_should_partition_fns,
60
+ )
61
+ from torch._inductor.utils import (
62
+ _unstable_customized_partition_wrapper,
63
+ is_cudagraph_unsafe_op,
64
+ maybe_log_cudagraph_partition,
65
+ )
66
+
67
+ # Allow users to manually specify if a node should be partitioned
68
+ # Can only do this for FallbackKernels
69
+ ir_node = node.node
70
+ if isinstance(ir_node, ir.FallbackKernel):
71
+ operator = ir_node.op_overload
72
+ if operator is not None and operator in _custom_should_partition_fns:
73
+ return True
74
+
75
+ # When not using cudagraphs, keep all kernels in the `call` function
76
+ # instead of graph partition functions, since graph partition only brings
77
+ # benefit to cudagraph
78
+ if (
79
+ not torch._inductor.config.triton.cudagraphs
80
+ and _unstable_customized_partition_wrapper.wrapper is None
81
+ ):
82
+ return True
83
+
84
+ # avoid duplicating logs when should_partition is called multiple times
85
+ # on the same node
86
+ def noop_log(msg: str, node: BaseSchedulerNode | None) -> None:
87
+ return
88
+
89
+ log_partition_reason = maybe_log_cudagraph_partition if should_log else noop_log
90
+
91
+ if isinstance(node, FusedSchedulerNode):
92
+ return any(self.should_partition(snode) for snode in node.snodes)
93
+
94
+ assert node.node is not None
95
+
96
+ if not node.is_gpu():
97
+ log_partition_reason("non gpu ops", node=node)
98
+
99
+ return True
100
+
101
+ if isinstance(node.node, ir.DeviceCopy):
102
+ log_partition_reason("DeviceCopy ops", node=node)
103
+ return True
104
+
105
+ if isinstance(node.node, ir.Conditional):
106
+ log_partition_reason("Conditional ops", node=node)
107
+ return True
108
+
109
+ if getattr(node.node, "unbacked_bindings", None):
110
+ log_partition_reason("unbacked binding ops", node=node)
111
+ return True
112
+
113
+ if is_cudagraph_unsafe_op(node.node):
114
+ log_partition_reason("CUDAGraph-unsafe custom ops", node=node)
115
+ return True
116
+
117
+ return False
118
+
119
+
120
+def _update_scheduler_patched(self) -> None:
121
+ # Copied from torch._inductor.graph.GrahLowering._update_scheduler. Patches
122
+ # this method so that we can patch Scheduler.should_partition with the
123
+ # function above
124
+ """
125
+ (Re)initializes the scheduler member. When initializing the scheduler, no CUBIN
126
+ files should be generated (to avoid biasing any benchmarks and pessimizing
127
+ fusion decisions).
128
+ """
129
+ import torch._inductor.config as config
130
+ from torch._inductor.scheduler import Scheduler
131
+
132
+ Scheduler.should_partition = should_partition_patched
133
+
134
+ with config.patch("triton.store_cubin", False):
135
+ self.scheduler = Scheduler(self.operations)
136
+
137
+
138
+if version.parse(str(torch.__version__)) == version.parse("2.9.0"):
139
+ from torch._inductor.graph import GraphLowering
140
+
141
+ GraphLowering._update_scheduler = _update_scheduler_patched