| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | from collections.abc import Iterable |
| 5 | |
| 6 | import torch.fx |
| 7 | from torch import SymInt |
| 8 | from torch.fx.experimental.symbolic_shapes import statically_known_true |
| 9 | |
| 10 | from vllm.logger import init_logger |
| 11 | |
| 12 | from ..fx_utils import is_func |
| 13 | from ..vllm_inductor_pass import VllmInductorPass |
| 14 | |
| 15 | logger = init_logger(__name__) |
| 16 | |
| 17 | |
| 18 | class NoOpEliminationPass(VllmInductorPass): |
| 19 | """This is an inductor pass that removes redundant reshape/slice operations. |
| 20 | It is required for RMSNorm-quant fusion to work properly. |
| 21 | That's because apply_fp8_linear adds a reshape, which is redundant |
| 22 | in the 2D-case. Additionally, torch internal no-op elimination pass does |
| 23 | not handle certain slice variants. |
| 24 | |
| 25 | Cases handled: |
| 26 | 1. A chain of reshapes is equivalent to the last reshape called on the |
| 27 | base tensor (input of the first reshape). |
| 28 | 2. A reshape that produces the shape of the input is redundant |
| 29 | 3. A slice that produces the shape of the input is redundant |
| 30 | |
| 31 | Example graph 1: |
| 32 | mul_1: "f16[s0, 4096]" = ... |
| 33 | view_1: "f16[s0, 128, 32]" = torch.reshape(mul_1, [-1, 128, 32]) |
| 34 | view_2: "f16[s0, 4096]" = torch.reshape(view_2, [-1, 4096]) |
| 35 | view_3: "f16[s0, 128, 32]" = torch.reshape(view_3, [-1, 128, 32]) |
| 36 | |
| 37 | Can be replaced with: |
| 38 | mul_1: "f16[s0, 4096]" = ... |
| 39 | view_3: "f16[s0, 128, 32]" = ... |
| 40 | |
| 41 | Example graph 2: |
| 42 | getitem_1: "f16[s0, 4096]" = ... |
| 43 | view_1: "f16[s0, 4096]" = torch.reshape(getitem_1, [-1, 4096]) |
| 44 | at = auto_functionalized(static_scaled_fp8_quant, input = view_1, ...) |
| 45 | out: "f8e4m3fn[s0, 4096]" = at[1] |
| 46 | |
| 47 | Can be replaced with: |
| 48 | getitem_1: "f16[s0, 4096]" = ... |
| 49 | at = auto_functionalized(static_scaled_fp8_quant, input = getitem_1, ...) |
| 50 | out: "f8e4m3fn[s0, 4096]" = at[1] |
| 51 | |
| 52 | Example graph 3: |
| 53 | arg0: "s0" = SymInt(s0) |
| 54 | scaled_mm: "f16[s0, 4096]" = ... |
| 55 | slice_1: "f16[s0, 4096]" = torch.slice(scaled_mm, -1, 0, arg0) |
| 56 | at = auto_functionalized(fused_add_rms_norm, input = slice_1, ...) |
| 57 | out: "f16[s0, 4096]" = torch.slice_scatter(scaled_mm, at[1], 0, 0, arg0) |
| 58 | |
| 59 | Can be replaced with: |
| 60 | arg0: "s0" = SymInt(s0) |
| 61 | scaled_mm: "f16[s0, 4096]" = ... |
| 62 | at = auto_functionalized(fused_add_rms_norm, input = scaled_mm, ...) |
| 63 | out: "f16[s0, 4096]" = at[1] |
| 64 | """ |
| 65 | |
| 66 | @VllmInductorPass.time_and_log |
| 67 | def __call__(self, graph: torch.fx.Graph) -> None: |
| 68 | count = 0 |
| 69 | # Remove no-op reshapes/views: |
| 70 | for node in graph.nodes: |
| 71 | if is_func(node, torch.ops.aten.reshape.default): |
| 72 | # Case 1: rewrite reshape chains to reshapes on the base tensor |
| 73 | input = node.args[0] |
| 74 | # If the input is a reshape, rebind to that node |
| 75 | if is_func(input, torch.ops.aten.reshape.default): |
| 76 | # The new input is guaranteed not to be a reshape, |
| 77 | # because we process nodes in order |
| 78 | node.update_arg(0, input.args[0]) |
| 79 | if len(input.users) == 0: |
| 80 | graph.erase_node(input) |
| 81 | count += 1 |
| 82 | |
| 83 | # remove reshape/slice if it produces the original shape |
| 84 | if is_func(node, torch.ops.aten.reshape.default) or is_func( |
| 85 | node, torch.ops.aten.slice.Tensor |
| 86 | ): |
| 87 | input = node.args[0] |
| 88 | input_shape = input.meta["val"].shape |
| 89 | output_shape = node.meta["val"].shape |
| 90 | if self.all_dims_equivalent(input_shape, output_shape): |
| 91 | node.replace_all_uses_with(input) |
| 92 | graph.erase_node(node) |
| 93 | count += 1 |
| 94 | elif is_func(node, torch.ops.aten.slice_scatter.default): |
| 95 | base, view, dim_index, start, end = node.args[:5] |
| 96 | base_shape = base.meta["val"].shape |
| 97 | view_shape = view.meta["val"].shape |
| 98 | |
| 99 | if self.all_dims_equivalent(base_shape, view_shape): |
| 100 | node.replace_all_uses_with(view) |
| 101 | graph.erase_node(node) |
| 102 | count += 1 |
| 103 | |
| 104 | logger.debug("Removed %s no-op reshapes and slices", count) |
| 105 | |
| 106 | # ---------------------- Shape comparison helpers ---------------------- |
| 107 | def dims_equivalent(self, dim: int | SymInt, i_dim: int | SymInt) -> bool: |
| 108 | """This function checks if two dimensions are equivalent. |
| 109 | |
| 110 | Args: |
| 111 | dim: The dimension arg to reshape/slice |
| 112 | i_dim: The corresponding dimension in the input tensor |
| 113 | |
| 114 | Returns: |
| 115 | Are the dimensions equivalent? |
| 116 | |
| 117 | There are two cases in which the dimensions are equivalent: |
| 118 | 1. The dimensions are equal (both integers) |
| 119 | 2. The dimensions both correspond to the same SymInt |
| 120 | |
| 121 | """ |
| 122 | # Case 1 |
| 123 | return statically_known_true(dim == i_dim) # type: ignore[no-any-return] |
| 124 | |
| 125 | def all_dims_equivalent( |
| 126 | self, dims: Iterable[int | SymInt], i_dims: Iterable[int | SymInt] |
| 127 | ) -> bool: |
| 128 | dims_ = list(dims) |
| 129 | i_dims_ = list(i_dims) |
| 130 | if len(dims_) != len(i_dims_): |
| 131 | # Different ranks can't be equivalent |
| 132 | return False |
| 133 | return all(self.dims_equivalent(s, i_s) for s, i_s in zip(dims, i_dims)) |