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
[torch.compile] Make inductor partition rules respect splitting_ops #25691 (#25845)
Signed-off-by: baonudesifeizhai <baonudesifeizhai@gmail.com> Signed-off-by: baonudesifeizhai <85092850+baonudesifeizhai@users.noreply.github.com> Co-authored-by: Luka Govedič <ProExpertProg@users.noreply.github.com>
baonudesifeizhai committed
Oct 10, 2025 at 12:35 UTC
cddce79fdae2bb843e8988d7f0cd145e918b2d1f
9 files changed
+267
-112
tests/compile/piecewise/test_multiple_graphs.py
+2
-2
@@ -198,7 +198,7 @@ def test_multi_graph_piecewise_compile_outputs_equal():
198
compilation_config=CompilationConfig(
199
level=CompilationLevel.PIECEWISE,
200
use_cudagraph=True,
201
- splitting_ops=["silly.attention"],
201
+ splitting_ops=["silly::attention"],
202
cudagraph_capture_sizes=[1, 2],
203
)
204
)
@@ -267,7 +267,7 @@ def test_multi_graph_piecewise_compile_outputs_equal():
267
compilation_config=CompilationConfig(
268
level=CompilationLevel.PIECEWISE,
269
use_cudagraph=False,
270
- splitting_ops=["silly.attention"],
270
+ splitting_ops=["silly::attention"],
271
)
272
)
273
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
tests/compile/piecewise/test_simple.py
+2
-2
@@ -127,7 +127,7 @@ def _run_simple_model(
127
@torch.inference_mode()
128
def test_simple_piecewise_compile(use_inductor):
129
_run_simple_model(
130
- splitting_ops=["silly.attention"],
130
+ splitting_ops=["silly::attention"],
131
use_inductor_graph_partition=False,
132
use_inductor=use_inductor,
133
# 2 * num_layers + 1
@@ -142,7 +142,7 @@ def test_simple_piecewise_compile(use_inductor):
142
143
144
@torch.inference_mode()
145
-@pytest.mark.parametrize("splitting_ops", [["silly.attention"], []])
145
+@pytest.mark.parametrize("splitting_ops", [["silly::attention"], []])
146
def test_simple_inductor_graph_partition(splitting_ops, monkeypatch):
147
if not is_torch_equal_or_newer("2.9.0.dev"):
148
pytest.skip("inductor graph partition is only available in PyTorch 2.9+")
tests/compile/piecewise/test_toy_llama.py
+2
-2
@@ -268,7 +268,7 @@ def run_model(
268
cudagraph_capture_sizes=[1, 2],
269
)
270
if split_attn:
271
- compilation_config.splitting_ops = ["silly.attention"]
271
+ compilation_config.splitting_ops = ["silly::attention"]
272
cudagraph_runtime_mode = CUDAGraphMode.PIECEWISE
273
else:
274
compilation_config = CompilationConfig(
@@ -438,7 +438,7 @@ def benchmark():
438
compilation_config = CompilationConfig(
439
level=CompilationLevel.PIECEWISE,
440
use_cudagraph=True,
441
- splitting_ops=["silly.attention"],
441
+ splitting_ops=["silly::attention"],
442
cudagraph_capture_sizes=cudagraph_sizes,
443
)
444
else:
tests/compile/test_config.py
+53
-29
@@ -4,10 +4,12 @@ import pytest
4
5
from vllm.compilation.counter import compilation_counter
6
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
7
-from vllm.utils import _is_torch_equal_or_newer
7
+from vllm.config.compilation import CompilationLevel
8
+from vllm.utils import _is_torch_equal_or_newer, is_torch_equal_or_newer
9
10
11
def test_version():
12
+ # Test the version comparison logic using the private function
13
assert _is_torch_equal_or_newer("2.8.0.dev20250624+cu128", "2.8.0.dev")
14
assert _is_torch_equal_or_newer("2.8.0a0+gitc82a174", "2.8.0.dev")
15
assert _is_torch_equal_or_newer("2.8.0", "2.8.0.dev")
@@ -17,6 +19,9 @@ def test_version():
19
20
def test_use_cudagraphs_dynamic():
21
vllm_config = VllmConfig()
22
+ # Default V1 configuration now starts without cudagraphs enabled; the
23
+ # engine decides when to capture based on runtime settings instead of a
24
+ # blanket default.
25
assert vllm_config.compilation_config.use_cudagraph
26
27
@@ -137,58 +142,77 @@ def test_enforce_eager(vllm_runner, monkeypatch):
142
def test_splitting_ops_dynamic():
143
# Default config
144
config = VllmConfig()
140
- assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL_AND_PIECEWISE
141
- assert config.compilation_config.splitting_ops_contain_attention()
145
+ # Default V1 config leaves cudagraph mode unset; splitting ops are only
146
+ # populated when the engine decides to use piecewise compilation.
147
+ assert config.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
148
+ assert not config.compilation_config.splitting_ops_contain_attention()
149
150
# When use_inductor_graph_partition=True
144
- if _is_torch_equal_or_newer("2.9.0.dev"):
145
- # inductor graph partition is only available in PyTorch 2.9+.
146
- # this is a fast config check so we are not using pytest.skip.
151
+ if is_torch_equal_or_newer("2.9.0.dev"):
152
config = VllmConfig(
153
compilation_config=CompilationConfig(
149
- use_inductor_graph_partition=True, splitting_ops=["silly_attention"]
154
+ level=CompilationLevel.PIECEWISE,
155
+ use_inductor_graph_partition=True,
156
+ splitting_ops=["vllm::unified_attention"],
157
)
158
)
152
- # should ignore splitting_ops
153
- assert config.compilation_config.splitting_ops == []
159
+ # with inductor partition we use splitting_ops directly for
160
+ # partition rules
161
+ assert config.compilation_config.splitting_ops == ["vllm::unified_attention"]
162
155
- # When attn_fusion pass enabled.
163
+ # When attn_fusion pass enabled, splitting_ops now default to attention ops.
164
config = VllmConfig(
165
compilation_config=CompilationConfig(
166
+ level=CompilationLevel.PIECEWISE,
167
pass_config={"enable_attn_fusion": True, "enable_noop": True},
168
custom_ops=["+quant_fp8"],
169
cudagraph_mode=CUDAGraphMode.PIECEWISE,
170
)
171
)
163
- assert config.compilation_config.splitting_ops == []
164
- # cudagraph mode also fall back to FULL
165
- assert config.compilation_config.cudagraph_mode == CUDAGraphMode.FULL
166
-
167
- # splitting_ops can not contain attention ops when attn_fusion
168
- # pass enabled.
169
- with pytest.raises(AssertionError):
170
- config = VllmConfig(
171
- compilation_config=CompilationConfig(
172
- pass_config={"enable_attn_fusion": True, "enable_noop": True},
173
- custom_ops=["+quant_fp8"],
174
- cudagraph_mode=CUDAGraphMode.PIECEWISE,
175
- # work around for accessing all attntion ops
176
- splitting_ops=CompilationConfig()._attention_ops,
177
- )
178
- )
172
+ # With the new simplified logic, attention fusion works with splitting_ops
173
+ assert config.compilation_config.splitting_ops_contain_attention()
174
+ # cudagraph mode remains PIECEWISE
175
+ assert config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE
176
177
# When both use_inductor_graph_partition and attn_fusion pass enabled.
181
- if _is_torch_equal_or_newer("2.9.0.dev"):
178
+ if is_torch_equal_or_newer("2.9.0.dev"):
179
config = VllmConfig(
180
compilation_config=CompilationConfig(
181
+ level=CompilationLevel.PIECEWISE,
182
use_inductor_graph_partition=True,
183
pass_config={"enable_attn_fusion": True, "enable_noop": True},
184
custom_ops=["+quant_fp8"],
185
cudagraph_mode=CUDAGraphMode.PIECEWISE,
186
)
187
)
190
- assert config.compilation_config.splitting_ops == []
191
- # enable_attn_fusion is directly support under
188
+ # With inductor graph partition, attn_fusion and splitting_ops
189
+ # work together. Default splitting_ops include attention ops.
190
+ assert config.compilation_config.splitting_ops_contain_attention()
191
+ # enable_attn_fusion is directly supported under
192
# use_inductor_graph_partition=True, and cudagraph_mode
193
# is unchanged.
194
assert config.compilation_config.cudagraph_mode == CUDAGraphMode.PIECEWISE
195
+
196
+
197
+def test_resolve_operator_overload():
198
+ import torch
199
+
200
+ from vllm.compilation.partition_rules import resolve_defined_ops
201
+
202
+ # Test valid operator names
203
+ resolved = resolve_defined_ops(["aten::mm.default", "aten::addmm.default"])
204
+ assert len(resolved) == 2
205
+ assert resolved[0] is torch.ops.aten.mm.default
206
+ assert resolved[1] is torch.ops.aten.addmm.default
207
+
208
+ # Test that invalid operators are skipped (not raising exceptions)
209
+ resolved = resolve_defined_ops(
210
+ [
211
+ "aten::mm.default",
212
+ "aten::nonexistent_op.default", # This should be skipped
213
+ "aten::addmm.default",
214
+ ]
215
+ )
216
+ assert len(resolved) == 2 # Only 2 valid ops
217
+ assert resolved[0] is torch.ops.aten.mm.default
218
+ assert resolved[1] is torch.ops.aten.addmm.default
tests/compile/test_decorator.py
+3
-3
@@ -71,7 +71,7 @@ def test_ignore_torch_compile_decorator():
71
compilation_config=CompilationConfig(
72
level=CompilationLevel.PIECEWISE,
73
use_cudagraph=True,
74
- splitting_ops=["silly.attention"],
74
+ splitting_ops=["silly::attention"],
75
cudagraph_capture_sizes=[1, 2],
76
)
77
)
@@ -186,7 +186,7 @@ def test_conditional_compile_enable_if():
186
compilation_config=CompilationConfig(
187
level=CompilationLevel.PIECEWISE,
188
use_cudagraph=True,
189
- splitting_ops=["silly.attention"],
189
+ splitting_ops=["silly::attention"],
190
cudagraph_capture_sizes=[1, 2],
191
),
192
)
@@ -218,7 +218,7 @@ def test_conditional_compile_enable_if():
218
compilation_config=CompilationConfig(
219
level=CompilationLevel.PIECEWISE,
220
use_cudagraph=True,
221
- splitting_ops=["silly.attention"],
221
+ splitting_ops=["silly::attention"],
222
cudagraph_capture_sizes=[1, 2],
223
),
224
)
vllm/compilation/backends.py
+44
-8
@@ -15,6 +15,11 @@ import torch.fx as fx
15
from torch._dispatch.python import enable_python_dispatcher
16
17
import vllm.envs as envs
18
+from vllm.compilation.inductor_pass import pass_context
19
+from vllm.compilation.partition_rules import (
20
+ inductor_partition_rule_context,
21
+ resolve_defined_ops,
22
+)
23
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
24
from vllm.logger import init_logger
25
from vllm.platforms import current_platform
@@ -76,6 +81,21 @@ class CompilerManager:
81
def compute_hash(self, vllm_config: VllmConfig) -> str:
82
return self.compiler.compute_hash(vllm_config)
83
84
+ @contextmanager
85
+ def compile_context(self, runtime_shape: Optional[int] = None):
86
+ """Provide compilation context for the duration of compilation to set
87
+ any torch global properties we want to scope to a single Inductor
88
+ compilation (e.g. partition rules, pass context)."""
89
+ with pass_context(runtime_shape):
90
+ if self.compilation_config.use_inductor_graph_partition:
91
+ inductor_partition_ops = resolve_defined_ops(
92
+ self.compilation_config.splitting_ops
93
+ )
94
+ with inductor_partition_rule_context(inductor_partition_ops):
95
+ yield
96
+ else:
97
+ yield
98
+
99
def initialize_cache(
100
self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
101
):
@@ -197,9 +217,15 @@ class CompilerManager:
217
maybe_key = None
218
else:
219
maybe_key = f"artifact_shape_{runtime_shape}_subgraph_{graph_index}"
200
- compiled_graph, handle = self.compiler.compile(
201
- graph, example_inputs, additional_inductor_config, runtime_shape, maybe_key
202
- )
220
+
221
+ with self.compile_context(runtime_shape):
222
+ compiled_graph, handle = self.compiler.compile(
223
+ graph,
224
+ example_inputs,
225
+ additional_inductor_config,
226
+ runtime_shape,
227
+ maybe_key,
228
+ )
229
230
assert compiled_graph is not None, "Failed to compile the graph"
231
@@ -258,7 +284,7 @@ class SplitItem:
284
285
286
def split_graph(
261
- graph: fx.GraphModule, ops: list[str]
287
+ graph: fx.GraphModule, resolved_ops: list[torch._ops.OpOverload]
288
) -> tuple[fx.GraphModule, list[SplitItem]]:
289
# split graph by ops
290
subgraph_id = 0
@@ -267,7 +293,12 @@ def split_graph(
293
for node in graph.graph.nodes:
294
if node.op in ("output", "placeholder"):
295
continue
270
- if node.op == "call_function" and str(node.target) in ops:
296
+ # Match node.target against resolved_ops
297
+ # node.target can be OpOverloadPacket, need to check .default
298
+ if node.op == "call_function" and (
299
+ node.target in resolved_ops
300
+ or (hasattr(node.target, "default") and node.target.default in resolved_ops)
301
+ ):
302
subgraph_id += 1
303
node_to_subgraph_id[node] = subgraph_id
304
split_op_graphs.append(subgraph_id)
@@ -615,9 +646,14 @@ class VllmBackend:
646
self.graph = graph
647
self.configure_post_pass()
648
618
- self.split_gm, self.piecewise_graphs = split_graph(
619
- graph, self.compilation_config.splitting_ops
620
- )
649
+ if self.compilation_config.use_inductor_graph_partition:
650
+ # Let Inductor decide partitioning; avoid FX-level pre-splitting.
651
+ fx_split_ops: list[str] = []
652
+ else:
653
+ fx_split_ops = self.compilation_config.splitting_ops or []
654
+
655
+ resolved_split_ops = resolve_defined_ops(fx_split_ops)
656
+ self.split_gm, self.piecewise_graphs = split_graph(graph, resolved_split_ops)
657
658
from torch._dynamo.utils import lazy_format_graph_code
659
vllm/compilation/compiler_interface.py
+12
-16
@@ -17,8 +17,6 @@ from vllm.compilation.counter import compilation_counter
17
from vllm.config import VllmConfig
18
from vllm.utils import is_torch_equal_or_newer
19
20
-from .inductor_pass import pass_context
21
-
20
21
class CompilerInterface:
22
"""
@@ -209,13 +207,12 @@ class InductorStandaloneAdaptor(CompilerInterface):
207
208
from torch._inductor import standalone_compile
209
212
- with pass_context(runtime_shape):
213
- compiled_graph = standalone_compile(
214
- graph,
215
- example_inputs,
216
- dynamic_shapes=dynamic_shapes,
217
- options={"config_patches": current_config},
218
- )
210
+ compiled_graph = standalone_compile(
211
+ graph,
212
+ example_inputs,
213
+ dynamic_shapes=dynamic_shapes,
214
+ options={"config_patches": current_config},
215
+ )
216
217
# Save the compiled artifact to disk in the specified path
218
assert key is not None
@@ -462,13 +459,12 @@ class InductorAdaptor(CompilerInterface):
459
torch._functorch.config.patch(enable_remote_autograd_cache=False)
460
)
461
465
- with pass_context(runtime_shape):
466
- compiled_graph = compile_fx(
467
- graph,
468
- example_inputs,
469
- inner_compile=hijacked_compile_fx_inner,
470
- config_patches=current_config,
471
- )
462
+ compiled_graph = compile_fx(
463
+ graph,
464
+ example_inputs,
465
+ inner_compile=hijacked_compile_fx_inner,
466
+ config_patches=current_config,
467
+ )
468
469
# We treat VLLM_DISABLE_COMPILE_CACHE as the overall switch for torch
470
# compilation cache. So turn off the checks if we disable the
vllm/compilation/partition_rules.py
new
+95
@@ -0,0 +1,95 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+
4
+from __future__ import annotations
5
+
6
+import contextlib
7
+from typing import TYPE_CHECKING
8
+
9
+from torch._library.utils import lookup_op
10
+
11
+from vllm.logger import init_logger
12
+
13
+if TYPE_CHECKING:
14
+ import torch
15
+
16
+logger = init_logger(__name__)
17
+
18
+
19
+def resolve_defined_ops(op_names: list[str]) -> list[torch._ops.OpOverload]:
20
+ """Resolve operator names to OpOverload objects.
21
+
22
+ Skips operators that fail to resolve (e.g., operators not registered or
23
+ model-specific operators not present in the current model).
24
+
25
+ Note: Users should inspect the operator graph before lowering and ensure
26
+ the specified operators are present in the final graph. Built-in PyTorch
27
+ operators (aten::*, torch::*) may be decomposed, fused, or transformed
28
+ during Inductor's compilation passes, so use them with caution.
29
+
30
+ Args:
31
+ op_names: List of operator names in PyTorch format
32
+ (e.g., "vllm::unified_attention")
33
+
34
+ Returns:
35
+ List of successfully resolved operator overloads
36
+ """
37
+ resolved = []
38
+ for op_name in op_names:
39
+ try:
40
+ resolved.append(lookup_op(op_name))
41
+ except Exception:
42
+ # Skip operators that don't exist (e.g., model-specific ops)
43
+ logger.warning(
44
+ "Failed to resolve operator for Inductor partition: %s", op_name
45
+ )
46
+ continue
47
+
48
+ return resolved
49
+
50
+
51
+@contextlib.contextmanager
52
+def inductor_partition_rule_context(overloads: list[torch._ops.OpOverload]):
53
+ """Context manager to temporarily register Inductor partition rules.
54
+
55
+ Registers custom partition rules for specified operators, forcing the
56
+ Inductor scheduler to partition the graph at these operators. The rules
57
+ are automatically restored to their previous state on exit.
58
+
59
+ Note: Callers should use resolve_defined_ops() to convert operator names
60
+ to OpOverload objects before calling this function.
61
+
62
+ Args:
63
+ overloads: List of resolved operator overload objects.
64
+ """
65
+ if not overloads:
66
+ logger.debug("No partition ops provided; skipping rule registration.")
67
+ yield
68
+ return
69
+
70
+ from torch._inductor.scheduler import ( # type: ignore
71
+ _custom_should_partition_fns,
72
+ register_should_partition_rule,
73
+ )
74
+
75
+ def _always_partition(*_args, **_kwargs):
76
+ return True
77
+
78
+ # Save current state before registering
79
+ saved_rules = _custom_should_partition_fns.copy()
80
+
81
+ for overload in overloads:
82
+ register_should_partition_rule(
83
+ overload,
84
+ _always_partition,
85
+ )
86
+
87
+ logger.debug("Registered inductor partition rules for %d operators", len(overloads))
88
+
89
+ try:
90
+ yield
91
+ finally:
92
+ # Clear and restore previous state
93
+ _custom_should_partition_fns.clear()
94
+ _custom_should_partition_fns.update(saved_rules)
95
+ logger.debug("Restored previous partition rules state.")
vllm/config/compilation.py
+54
-50
@@ -209,8 +209,23 @@ class CompilationConfig:
209
disabled when running with Inductor: level>=PIECEWISE and use_inductor=True.
210
Inductor generates (fused) Triton kernels for disabled custom ops."""
211
splitting_ops: Optional[list[str]] = None
212
- """A list of ops to split the full graph into subgraphs, used in piecewise
213
- compilation."""
212
+ """A list of ops to exclude from cudagraphs, used in piecewise compilation.
213
+
214
+ The behavior depends on use_inductor_graph_partition:
215
+
216
+ - When use_inductor_graph_partition=False (default):
217
+ These ops are used for Dynamo FX-level graph splitting. The graph is
218
+ split at these ops before Inductor compilation, creating separate
219
+ subgraphs for cudagraph capture.
220
+
221
+ - When use_inductor_graph_partition=True:
222
+ These ops are used to register Inductor partition rules. The graph
223
+ partitioning happens at Inductor codegen time after all passes and
224
+ fusions are finished, allowing compilation and custom passes to operate
225
+ on the full graph while still excluding these ops from cudagraphs.
226
+
227
+ If None, defaults to attention ops for piecewise cudagraphs.
228
+ If empty list [], no ops are excluded (suitable for full cudagraphs)."""
229
230
# Inductor capture
231
use_inductor: bool = True
@@ -367,18 +382,19 @@ class CompilationConfig:
382
model code, e.g., Attention, FusedMOE when dp_size>1."""
383
384
# Attention ops; used for piecewise cudagraphs
385
+ # Use PyTorch operator format: "namespace::name"
386
_attention_ops: ClassVar[list[str]] = [
371
- "vllm.unified_attention",
372
- "vllm.unified_attention_with_output",
373
- "vllm.unified_mla_attention",
374
- "vllm.unified_mla_attention_with_output",
375
- "vllm.mamba_mixer2",
376
- "vllm.mamba_mixer",
377
- "vllm.short_conv",
378
- "vllm.linear_attention",
379
- "vllm.plamo2_mamba_mixer",
380
- "vllm.gdn_attention",
381
- "vllm.sparse_attn_indexer",
387
+ "vllm::unified_attention",
388
+ "vllm::unified_attention_with_output",
389
+ "vllm::unified_mla_attention",
390
+ "vllm::unified_mla_attention_with_output",
391
+ "vllm::mamba_mixer2",
392
+ "vllm::mamba_mixer",
393
+ "vllm::short_conv",
394
+ "vllm::linear_attention",
395
+ "vllm::plamo2_mamba_mixer",
396
+ "vllm::gdn_attention",
397
+ "vllm::sparse_attn_indexer",
398
]
399
400
def compute_hash(self) -> str:
@@ -654,31 +670,25 @@ class CompilationConfig:
670
671
def set_splitting_ops_for_inductor_graph_partition(self):
672
assert self.use_inductor_graph_partition
657
- use_inductor_graph_partition_msg = (
658
- "When use_inductor_graph_partition=True, splitting_ops "
659
- "are ignored and set to an empty list. Instead, "
660
- '"tags=(torch._C.Tag.cudagraph_unsafe, )," is '
661
- "used to annotate custom ops for graph partition."
662
- )
663
- if self.splitting_ops is not None and len(self.splitting_ops) > 0:
664
- logger.warning_once(use_inductor_graph_partition_msg)
665
- self.splitting_ops = []
673
+ if self.splitting_ops is None:
674
+ self.splitting_ops = list(self._attention_ops)
675
676
def set_splitting_ops_for_attn_fusion(self):
677
assert self.pass_config.enable_attn_fusion
669
- if self.splitting_ops is None:
670
- self.splitting_ops = []
671
- if self.cudagraph_mode.has_piecewise_cudagraphs():
672
- logger.warning_once(
673
- "enable_attn_fusion is incompatible with piecewise "
674
- "cudagraph when use_inductor_graph_partition is off."
675
- "In this case, splitting_ops will be set to empty "
676
- "list, and cudagraph_mode will be set to FULL. "
677
- "Please ensure you are using attention backends that "
678
- "support cudagraph or set cudagraph_mode to NONE "
679
- "explicitly if encountering any problems."
680
- )
681
- self.cudagraph_mode = CUDAGraphMode.FULL
678
+ # For dynamo-partition (non-inductor) attention fusion,
679
+ # set splitting_ops to empty to avoid splitting at attention ops
680
+ self.splitting_ops = []
681
+ if self.cudagraph_mode.has_piecewise_cudagraphs():
682
+ logger.warning_once(
683
+ "enable_attn_fusion is incompatible with piecewise "
684
+ "cudagraph when use_inductor_graph_partition is off. "
685
+ "In this case, splitting_ops will be set to empty "
686
+ "list, and cudagraph_mode will be set to FULL. "
687
+ "Please ensure you are using attention backends that "
688
+ "support cudagraph or set cudagraph_mode to NONE "
689
+ "explicitly if encountering any problems."
690
+ )
691
+ self.cudagraph_mode = CUDAGraphMode.FULL
692
693
assert not self.splitting_ops_contain_attention(), (
694
"attention ops should not be in splitting_ops "
@@ -691,23 +701,17 @@ class CompilationConfig:
701
)
702
703
def is_attention_compiled_piecewise(self) -> bool:
694
- use_fx_graph_piecewise_compilation = (
695
- self.level == CompilationLevel.PIECEWISE
696
- and self.splitting_ops_contain_attention()
697
- )
704
+ if not self.splitting_ops_contain_attention():
705
+ return False
706
699
- inductor_used = (
700
- self.level == CompilationLevel.PIECEWISE and self.use_inductor
701
- ) or (
702
- self.level >= CompilationLevel.DYNAMO_AS_IS and self.backend == "inductor"
703
- )
704
- use_inductor_piecewise_compilation = (
705
- inductor_used
706
- and self.use_inductor_graph_partition
707
- and not self.splitting_ops_contain_attention()
708
- )
707
+ if not self.use_inductor_graph_partition:
708
+ # Dynamo-level FX split case
709
+ return self.level == CompilationLevel.PIECEWISE
710
710
- return use_fx_graph_piecewise_compilation or use_inductor_piecewise_compilation
711
+ # Inductor partition case
712
+ return (
713
+ self.level > CompilationLevel.NO_COMPILATION and self.backend == "inductor"
714
+ )
715
716
def custom_op_log_check(self):
717
"""