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 13, 19:30 UTC
[Graph Partition][Cache] Use inductor partition ops config (#27702)
Signed-off-by: Boyuan Feng <boyuan@meta.com>
Boyuan Feng committed
Nov 5, 2025 at 05:04 UTC
6ab183813caa26813cf71de4991178dd0017164f
4 files changed
+37
-63
vllm/compilation/backends.py
+2
-3
@@ -97,10 +97,9 @@ class CompilerManager:
97
compilation (e.g. partition rules, pass context)."""
98
with pass_context(runtime_shape):
99
if self.compilation_config.use_inductor_graph_partition:
100
- inductor_partition_ops = resolve_defined_ops(
100
+ with inductor_partition_rule_context(
101
self.compilation_config.splitting_ops
102
- )
103
- with inductor_partition_rule_context(inductor_partition_ops):
102
+ ):
103
yield
104
else:
105
yield
vllm/compilation/partition_rules.py
+12
-27
@@ -3,15 +3,12 @@
3
4
import contextlib
5
import logging
6
-from typing import TYPE_CHECKING
6
7
+import torch
8
from torch._library.utils import lookup_op
9
10
from vllm.logger import init_logger
11
12
-if TYPE_CHECKING:
13
- import torch
14
-
12
logger = init_logger(__name__)
13
14
@@ -56,47 +53,35 @@ def resolve_defined_ops(op_names: list[str]) -> list["torch._ops.OpOverload"]:
53
54
55
@contextlib.contextmanager
59
-def inductor_partition_rule_context(overloads: list["torch._ops.OpOverload"]):
56
+def inductor_partition_rule_context(splitting_ops: list[str]):
57
"""Context manager to temporarily register Inductor partition rules.
58
59
Registers custom partition rules for specified operators, forcing the
60
Inductor scheduler to partition the graph at these operators. The rules
61
are automatically restored to their previous state on exit.
62
66
- Note: Callers should use resolve_defined_ops() to convert operator names
67
- to OpOverload objects before calling this function.
68
-
63
Args:
70
- overloads: List of resolved operator overload objects.
64
+ splitting_ops: List of operator names to partition on.
65
"""
72
- if not overloads:
66
+ if not splitting_ops:
67
logger.debug("No partition ops provided; skipping rule registration.")
68
yield
69
return
70
77
- from torch._inductor.scheduler import ( # type: ignore
78
- _custom_should_partition_fns,
79
- register_should_partition_rule,
80
- )
81
-
82
- def _always_partition(*_args, **_kwargs):
83
- return True
84
-
71
# Save current state before registering
86
- saved_rules = _custom_should_partition_fns.copy()
72
88
- for overload in overloads:
89
- register_should_partition_rule(
90
- overload,
91
- _always_partition,
92
- )
73
+ saved_splitting_ops: list[str] = list(
74
+ torch._inductor.config.custom_should_partition_ops
75
+ )
76
+ torch._inductor.config.custom_should_partition_ops = splitting_ops
77
94
- logger.debug("Registered inductor partition rules for %d operators", len(overloads))
78
+ logger.debug(
79
+ "Registered inductor partition rules for %d operators", len(splitting_ops)
80
+ )
81
82
try:
83
yield
84
finally:
85
# Clear and restore previous state
100
- _custom_should_partition_fns.clear()
101
- _custom_should_partition_fns.update(saved_rules)
86
+ torch._inductor.config.custom_should_partition_ops = saved_splitting_ops
87
logger.debug("Restored previous partition rules state.")
vllm/compilation/pass_manager.py
+1
-29
@@ -113,27 +113,6 @@ class PostGradPassManager(CustomGraphPass):
113
self.post_cleanup = PostCleanupPass(config)
114
self.fix_functionalization = FixFunctionalizationPass(config)
115
116
- # [HACK: Bug with Inductor graph partition and torch.compile cache]
117
- # In PyTorch 2.9, torch.compile has a bug where the graph
118
- # partition is not taken into account during caching.
119
- # Because vLLM's Mode.VLLM_COMPILE is the only mode that uses
120
- # Inductor graph partition, and VLLM_COMPILE implies there
121
- # is a PostGradPassManager, we put the list of operators to graph
122
- # partition into the PostGradPassManager's uuid (which
123
- # then gets incorporated into Inductor's FX graph cache key).
124
- # Remove this hack whenever torch.compile fixes it.
125
-
126
- # This is the list of operators that vLLM asks Inductor to split.
127
- self.inductor_splitting_ops = []
128
- if (
129
- config.compilation_config.use_inductor_graph_partition
130
- and config.compilation_config.splitting_ops is not None
131
- ):
132
- # Sort them so we're not dependent on the ordering.
133
- self.inductor_splitting_ops = sorted(
134
- config.compilation_config.splitting_ops
135
- )
136
-
116
def add(self, pass_: InductorPass):
117
assert isinstance(pass_, InductorPass)
118
self.passes.append(pass_)
@@ -144,16 +123,9 @@ class PostGradPassManager(CustomGraphPass):
123
affects compilation caching. Its uuid depends on the UUIDs of all
124
dependent passes and the pass config. See InductorPass for more info.
125
"""
147
- state = {
148
- "pass_config": self.pass_config.uuid(),
149
- "passes": [],
150
- "inductor_splitting_ops": [],
151
- }
126
+ state = {"pass_config": self.pass_config.uuid(), "passes": []}
127
for pass_ in self.passes:
128
state["passes"].append(pass_.uuid())
129
state["passes"].append(self.fix_functionalization.uuid())
130
156
- # See [HACK: Bug with Inductor graph partition and torch.compile cache]
157
- state["inductor_splitting_ops"].extend(self.inductor_splitting_ops)
158
-
131
return InductorPass.hash_dict(state)
vllm/env_override.py
+22
-4
@@ -272,7 +272,6 @@ def should_partition_patched(self, node, should_log: bool = False) -> bool:
272
from torch._inductor.scheduler import (
273
BaseSchedulerNode,
274
FusedSchedulerNode,
275
- _custom_should_partition_fns,
275
)
276
from torch._inductor.utils import (
277
_unstable_customized_partition_wrapper,
@@ -283,9 +282,21 @@ def should_partition_patched(self, node, should_log: bool = False) -> bool:
282
# Allow users to manually specify if a node should be partitioned
283
# Can only do this for FallbackKernels
284
ir_node = node.node
286
- if isinstance(ir_node, ir.FallbackKernel):
287
- operator = ir_node.op_overload
288
- if operator is not None and operator in _custom_should_partition_fns:
285
+ if isinstance(ir_node, torch._inductor.ir.FallbackKernel) and (
286
+ op := ir_node.op_overload
287
+ ):
288
+ op_overload_packet_name = op.name()
289
+ op_overload_name = (
290
+ f"{op_overload_packet_name}.{op._overloadname}"
291
+ if isinstance(op, torch._ops.OpOverload)
292
+ else op_overload_packet_name
293
+ )
294
+ if (
295
+ op_overload_packet_name
296
+ in torch._inductor.config.custom_should_partition_ops
297
+ or op_overload_name in torch._inductor.config.custom_should_partition_ops
298
+ ):
299
+ assert isinstance(op, torch._ops.OpOverload)
300
return True
301
302
# When not using cudagraphs, keep all kernels in the `call` function
@@ -355,6 +366,13 @@ def _update_scheduler_patched(self) -> None:
366
if is_torch_equal("2.9.0"):
367
from torch._inductor.codegen.wrapper import PythonWrapperCodegen
368
from torch._inductor.graph import GraphLowering
369
+ from torch.utils._config_module import _Config, _ConfigEntry
370
+
371
+ # `custom_should_partition_ops` is a new config after 2.9.0. So this would
372
+ # not overwrite any user configs.
373
+ torch._inductor.config._config["custom_should_partition_ops"] = _ConfigEntry(
374
+ _Config(default=[])
375
+ )
376
377
PythonWrapperCodegen.memory_plan_reuse = memory_plan_reuse_patched
378
GraphLowering._update_scheduler = _update_scheduler_patched