| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | """Replace ``slice_scatter`` and ``split_with_sizes`` nodes with a single |
| 4 | assignment if there are no users for the inplace tensor written to by |
| 5 | the slice_scatter call. |
| 6 | |
| 7 | The inplace rotary_embedding custom op takes in mutable query and key inputs |
| 8 | that are split+getitem outputs of a single qkv tensor. |
| 9 | When functionalized, we fetch the rotated query and key from the functionalized op |
| 10 | using `getitem` calls. However, we also write to the qkv tensor inplace using a |
| 11 | `slice_scatter`, then split the inplace tensor to get the output tensors again. |
| 12 | Instead, if the inplace tensor has no subsequent users, we can just replace the |
| 13 | `slice_scatter` and `split_with_sizes` nodes with the `getitem` calls. |
| 14 | |
| 15 | This is already done in fix_functionalization::FixFunctionalizationPass, but |
| 16 | writing a custom pass for it before defunctionalization allows matching against the |
| 17 | qkv split+rotary_embedding subpattern as part of e.g. the RoPE+KVCache fusion pass. |
| 18 | """ |
| 19 | |
| 20 | import operator |
| 21 | |
| 22 | import torch |
| 23 | from torch import fx |
| 24 | from torch._higher_order_ops.auto_functionalize import auto_functionalized |
| 25 | |
| 26 | from vllm.logger import init_logger |
| 27 | |
| 28 | from ..fx_utils import is_func |
| 29 | from ..vllm_inductor_pass import VllmInductorPass |
| 30 | |
| 31 | logger = init_logger(__name__) |
| 32 | |
| 33 | |
| 34 | class ScatterSplitReplacementPass(VllmInductorPass): |
| 35 | """Replace getitem+slice_scatter+split nodes with a single getitem when |
| 36 | the inplace subtensor written to by the slice_scatter has no other users. |
| 37 | |
| 38 | Here's an example graph with q_size = 512, kv_size = 64: |
| 39 | split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1) |
| 40 | at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k)) |
| 41 | q = operator.getitem(at, 1) |
| 42 | k = operator.getitem(at, 2) |
| 43 | torch.ops.aten.slice_scatter.default(qkv, q, [0, 512], -1) |
| 44 | torch.ops.aten.slice_scatter.default(qkv, k, [512, 512 + 64], -1) |
| 45 | split_with_sizes_2 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1) |
| 46 | q = operator.getitem(split_with_sizes_2, 0) |
| 47 | k = operator.getitem(split_with_sizes_2, 1) |
| 48 | v = operator.getitem(split_with_sizes_2, 2) |
| 49 | |
| 50 | After this pass, this sequence of nodes is replaced with: |
| 51 | split_with_sizes_1 = torch.ops.aten.split_with_sizes.default(qkv, (512, 64, 64), -1) |
| 52 | at = auto_functionalized(torch.ops._C.rotary_embedding.default(positions, q, k)) |
| 53 | q = operator.getitem(at, 1) |
| 54 | k = operator.getitem(at, 2) |
| 55 | v = operator.getitem(split_with_sizes_1, 2) |
| 56 | """ |
| 57 | |
| 58 | @VllmInductorPass.time_and_log |
| 59 | def __call__(self, graph: fx.Graph) -> None: |
| 60 | count = 0 |
| 61 | |
| 62 | target_ops = [torch.ops._C.rotary_embedding.default] |
| 63 | if hasattr(torch.ops.vllm, "rocm_aiter_triton_rotary_embedding"): |
| 64 | target_ops.append(torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default) |
| 65 | |
| 66 | for node in graph.nodes: |
| 67 | if not is_func(node, auto_functionalized): |
| 68 | continue |
| 69 | |
| 70 | kwargs = node.kwargs |
| 71 | at_target = node.args[0] |
| 72 | |
| 73 | if at_target in target_ops: |
| 74 | query = kwargs["query"] |
| 75 | key = kwargs["key"] |
| 76 | getitem_nodes = {} |
| 77 | for user in node.users: |
| 78 | if is_func(user, operator.getitem): |
| 79 | getitem_nodes[user.args[1]] = user |
| 80 | |
| 81 | if ( |
| 82 | is_func(query, operator.getitem) |
| 83 | and is_func(key, operator.getitem) |
| 84 | and query.args[0] == key.args[0] |
| 85 | and is_func(query.args[0], torch.ops.aten.split_with_sizes.default) |
| 86 | and all( |
| 87 | is_func(user, torch.ops.aten.slice_scatter.default) |
| 88 | for getitem_node in getitem_nodes.values() |
| 89 | for user in getitem_node.users |
| 90 | ) |
| 91 | ): |
| 92 | # Pattern where query and key are slices of a qkv tensor. |
| 93 | # While functionalized, results at [1] and [2] are scattered |
| 94 | # back into qkv, then split again to get query and key. |
| 95 | # If the inplace tensor has no other users, we can replace |
| 96 | # the slice_scatter+split nodes with the original results. |
| 97 | for user in getitem_nodes[1].users: |
| 98 | slice_scatter_1_node = user |
| 99 | if not is_func( |
| 100 | slice_scatter_1_node, torch.ops.aten.slice_scatter.default |
| 101 | ): |
| 102 | continue |
| 103 | |
| 104 | for user in getitem_nodes[2].users: |
| 105 | slice_scatter_2_node = user |
| 106 | if not is_func( |
| 107 | slice_scatter_2_node, torch.ops.aten.slice_scatter.default |
| 108 | ): |
| 109 | continue |
| 110 | |
| 111 | for user in slice_scatter_2_node.users: |
| 112 | split_node = user |
| 113 | if not is_func(split_node, torch.ops.aten.split_with_sizes.default): |
| 114 | continue |
| 115 | |
| 116 | split_getitem_users = {} |
| 117 | for user in split_node.users: |
| 118 | if is_func(user, operator.getitem): |
| 119 | split_getitem_users[user.args[1]] = user |
| 120 | |
| 121 | # Replace query node |
| 122 | split_getitem_users[0].replace_all_uses_with(getitem_nodes[1]) |
| 123 | graph.erase_node(split_getitem_users[0]) |
| 124 | # Replace key node |
| 125 | split_getitem_users[1].replace_all_uses_with(getitem_nodes[2]) |
| 126 | graph.erase_node(split_getitem_users[1]) |
| 127 | # Redirect value node to original qkv tensor |
| 128 | split_getitem_users[2].replace_input_with(split_node, query.args[0]) |
| 129 | |
| 130 | # Erase unused nodes |
| 131 | graph.erase_node(split_node) |
| 132 | graph.erase_node(slice_scatter_2_node) |
| 133 | graph.erase_node(slice_scatter_1_node) |
| 134 | |
| 135 | count += 1 |
| 136 | |
| 137 | logger.debug("Eliminated %d slice_scatter+split nodes", count) |