| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import contextlib |
| 5 | from collections.abc import Generator |
| 6 | |
| 7 | import torch |
| 8 | |
| 9 | from vllm.logger import init_logger |
| 10 | |
| 11 | logger = init_logger(__name__) |
| 12 | |
| 13 | |
| 14 | def should_split(node: torch.fx.Node, splitting_ops: list[str]) -> bool: |
| 15 | """Check if a node should be split for dynamo graph partition. |
| 16 | It operates on dynamo graph, so the node.target can be anything. |
| 17 | We need to check and split only on OpOverload and OpOverloadPacket. |
| 18 | """ |
| 19 | if node.op != "call_function": |
| 20 | return False |
| 21 | |
| 22 | target = node.target |
| 23 | |
| 24 | if isinstance(target, torch._ops.OpOverloadPacket): |
| 25 | # Example: "aten::add" |
| 26 | return target._qualified_op_name in splitting_ops |
| 27 | |
| 28 | if isinstance(target, torch._ops.OpOverload): |
| 29 | # Example: "aten::add" |
| 30 | packet_name = target.name() |
| 31 | |
| 32 | # Example: "aten::add.default" |
| 33 | op_overload_name = f"{packet_name}.{target._overloadname}" |
| 34 | return op_overload_name in splitting_ops or packet_name in splitting_ops |
| 35 | |
| 36 | return False |
| 37 | |
| 38 | |
| 39 | @contextlib.contextmanager |
| 40 | def inductor_partition_rule_context( |
| 41 | splitting_ops: list[str] | None, |
| 42 | ) -> Generator[None, None, None]: |
| 43 | """Context manager to temporarily register Inductor partition rules. |
| 44 | |
| 45 | Registers custom partition rules for specified operators, forcing the |
| 46 | Inductor scheduler to partition the graph at these operators. The rules |
| 47 | are automatically restored to their previous state on exit. |
| 48 | |
| 49 | Args: |
| 50 | splitting_ops: List of operator names to partition on. |
| 51 | |
| 52 | """ |
| 53 | if not splitting_ops: |
| 54 | logger.debug("No partition ops provided; skipping rule registration.") |
| 55 | yield |
| 56 | return |
| 57 | |
| 58 | # Save current state before registering |
| 59 | |
| 60 | saved_splitting_ops: list[str] = list( |
| 61 | torch._inductor.config.custom_should_partition_ops |
| 62 | ) |
| 63 | torch._inductor.config.custom_should_partition_ops = splitting_ops |
| 64 | |
| 65 | logger.debug( |
| 66 | "Registered inductor partition rules for %d operators", len(splitting_ops) |
| 67 | ) |
| 68 | |
| 69 | try: |
| 70 | yield |
| 71 | finally: |
| 72 | # Clear and restore previous state |
| 73 | torch._inductor.config.custom_should_partition_ops = saved_splitting_ops |
| 74 | logger.debug("Restored previous partition rules state.") |