| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import torch |
| 4 | from torch import fx |
| 5 | from torch._higher_order_ops.auto_functionalize import auto_functionalized |
| 6 | from torch._higher_order_ops.triton_kernel_wrap import TritonKernelWrapperFunctional |
| 7 | from torch._ops import HigherOrderOperator, OpOverload |
| 8 | |
| 9 | from vllm.config import VllmConfig |
| 10 | from vllm.logger import init_logger |
| 11 | |
| 12 | from ..fx_utils import is_func |
| 13 | from ..inductor_pass import get_pass_context |
| 14 | from ..vllm_inductor_pass import VllmInductorPass |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | def clone_preserves_layout(node: fx.Node, original_node: fx.Node) -> bool: |
| 20 | node_val = node.meta.get("val") |
| 21 | original_val = original_node.meta.get("val") |
| 22 | if node_val is None or original_val is None: |
| 23 | return True |
| 24 | |
| 25 | try: |
| 26 | node_stride = tuple(node_val.stride()) |
| 27 | original_stride = tuple(original_val.stride()) |
| 28 | node_storage_offset = node_val.storage_offset() |
| 29 | original_storage_offset = original_val.storage_offset() |
| 30 | except (AttributeError, RuntimeError): |
| 31 | return True |
| 32 | |
| 33 | return ( |
| 34 | node_stride == original_stride |
| 35 | and node_storage_offset == original_storage_offset |
| 36 | ) |
| 37 | |
| 38 | |
| 39 | def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool: |
| 40 | if user.op == "output": |
| 41 | return False |
| 42 | |
| 43 | if is_func(user, auto_functionalized): |
| 44 | # While autofunc writes to the node, |
| 45 | # this is a follow-up use we're not interested in. |
| 46 | # It is also guaranteed to be the final use, |
| 47 | # as auto_functionalized returns the tensor back for follow-up use. |
| 48 | return False |
| 49 | elif user.op == "call_function" and isinstance(user.target, HigherOrderOperator): |
| 50 | # By default, be conservative, assume this could be a write |
| 51 | # (except functional HOPs) |
| 52 | return not isinstance(user.target, TritonKernelWrapperFunctional) |
| 53 | |
| 54 | assert isinstance(user.target, OpOverload), ( |
| 55 | f"{node=} {user=} {user.op=} {user.target=}" |
| 56 | ) |
| 57 | schema = user.target._schema |
| 58 | assert len(user.args) <= len(schema.arguments) |
| 59 | for i, arg in enumerate(user.args): |
| 60 | # Only interested in writes to node |
| 61 | if arg is not node: |
| 62 | continue |
| 63 | |
| 64 | # If not a write, next arg could be |
| 65 | if schema.arguments[i].is_write: |
| 66 | return True |
| 67 | |
| 68 | # No writes found |
| 69 | return False |
| 70 | |
| 71 | |
| 72 | class UnsafeCloneEliminationPass(VllmInductorPass): |
| 73 | """This pass removes clone nodes that are no longer needed after vLLM IR lowering. |
| 74 | It uses donated_input_ids to eliminate clones of donated graph inputs, preserving |
| 75 | contents of non-donated graph inputs. |
| 76 | |
| 77 | It is "unsafe" because it does not (yet) take aliasing into account. Solving |
| 78 | aliasing is an open problem, so this pass intends to support known vLLM cases |
| 79 | and not guarantee soundness on general graphs. In the future, this pass will likely |
| 80 | support basic forms of aliasing to handle simple views (e.g. qkv -> q,k,v). |
| 81 | """ |
| 82 | |
| 83 | def __init__(self, vllm_config: VllmConfig) -> None: |
| 84 | super().__init__(vllm_config) |
| 85 | |
| 86 | @VllmInductorPass.time_and_log |
| 87 | def __call__(self, graph: fx.Graph) -> None: |
| 88 | count = 0 |
| 89 | node_to_idx = {node: i for i, node in enumerate(graph.nodes)} |
| 90 | pass_context = get_pass_context() |
| 91 | donated_input_ids = pass_context.donated_input_ids |
| 92 | logger.debug("Donated input ids: %s", donated_input_ids) |
| 93 | |
| 94 | for node in graph.nodes: |
| 95 | if not is_func(node, torch.ops.aten.clone.default): |
| 96 | continue |
| 97 | |
| 98 | original_node = node.args[0] |
| 99 | assert isinstance(original_node, fx.Node) |
| 100 | |
| 101 | if not clone_preserves_layout(node, original_node): |
| 102 | logger.debug( |
| 103 | "Clone removal not possible, clone changes layout: " |
| 104 | "original_node=%s node=%s", |
| 105 | original_node, |
| 106 | node, |
| 107 | ) |
| 108 | continue |
| 109 | |
| 110 | # Clone needs to be preserved if node is getting written to and |
| 111 | # the old value is used again. |
| 112 | # This could only happen if an inplace implementation was lowered. |
| 113 | # Then node (the clone) will have one write. |
| 114 | # TODO(luka) hopefully this can be removed once we lower functional graphs. |
| 115 | write_idxs = [ |
| 116 | node_to_idx[u] for u in node.users if user_writes_to_node(u, node) |
| 117 | ] |
| 118 | assert len(write_idxs) in (0, 1) |
| 119 | if write_idxs: |
| 120 | # Check if a user of original_node occurs after a write |
| 121 | write_idx = write_idxs[0] |
| 122 | if any( |
| 123 | node_to_idx[orig_user] > write_idx |
| 124 | for orig_user in original_node.users |
| 125 | ): |
| 126 | logger.debug( |
| 127 | "Clone removal not possible, " |
| 128 | "original_node=%s used after mutation on node=%s", |
| 129 | original_node, |
| 130 | node, |
| 131 | ) |
| 132 | continue |
| 133 | |
| 134 | # Check if a node is a (non-donated) graph input |
| 135 | if ( |
| 136 | original_node.op == "placeholder" |
| 137 | and node_to_idx[original_node] not in donated_input_ids |
| 138 | ): |
| 139 | logger.debug( |
| 140 | "Graph input %s not donated, cannot eliminate its clone", |
| 141 | original_node, |
| 142 | ) |
| 143 | continue |
| 144 | |
| 145 | logger.debug( |
| 146 | "Node %s is a redundant clone node of %s, removing it", |
| 147 | node, |
| 148 | original_node, |
| 149 | ) |
| 150 | node.replace_all_uses_with(original_node) |
| 151 | graph.erase_node(node) |
| 152 | count += 1 |
| 153 | |
| 154 | logger.debug("CloneCleanupPass removed %d clone nodes", count) |