| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import operator |
| 5 | from collections.abc import Iterable |
| 6 | |
| 7 | import torch |
| 8 | from torch._higher_order_ops.auto_functionalize import auto_functionalized |
| 9 | |
| 10 | from vllm.logger import init_logger |
| 11 | from vllm.platforms import current_platform |
| 12 | |
| 13 | from ..fx_utils import is_func |
| 14 | from ..vllm_inductor_pass import VllmInductorPass |
| 15 | |
| 16 | logger = init_logger(__name__) |
| 17 | |
| 18 | |
| 19 | class FixFunctionalizationPass(VllmInductorPass): |
| 20 | """This pass defunctionalizes certain nodes to avoid redundant tensor copies. |
| 21 | After this pass, DCE (dead-code elimination) should never be run, |
| 22 | as de-functionalized nodes may appear as dead code. |
| 23 | |
| 24 | To add new nodes to defunctionalize, add to the if-elif chain in __call__. |
| 25 | """ |
| 26 | |
| 27 | @VllmInductorPass.time_and_log |
| 28 | def __call__(self, graph: torch.fx.Graph) -> None: |
| 29 | # XPU does not support auto-functionalization yet. |
| 30 | # Will enable this when switch to vllm-xpu-kernels. |
| 31 | if current_platform.is_xpu(): |
| 32 | logger.debug( |
| 33 | "XPU platform does not support fix functionalizationpass currently." |
| 34 | ) |
| 35 | return |
| 36 | |
| 37 | self.nodes_to_remove: list[torch.fx.Node] = [] |
| 38 | count = 0 |
| 39 | |
| 40 | rope_targets = [torch.ops._C.rotary_embedding.default] |
| 41 | |
| 42 | if hasattr(torch.ops.vllm, "rocm_aiter_triton_rotary_embedding"): |
| 43 | rope_targets.append( |
| 44 | torch.ops.vllm.rocm_aiter_triton_rotary_embedding.default |
| 45 | ) |
| 46 | |
| 47 | for node in graph.nodes: |
| 48 | if not is_func(node, auto_functionalized): |
| 49 | continue # Avoid deep if-elif nesting |
| 50 | |
| 51 | kwargs = node.kwargs |
| 52 | at_target = node.args[0] |
| 53 | |
| 54 | if at_target in rope_targets: |
| 55 | query = kwargs["query"] |
| 56 | key = kwargs["key"] |
| 57 | getitem_nodes = self.getitem_users(node) |
| 58 | |
| 59 | if ( |
| 60 | is_func(query, operator.getitem) |
| 61 | and is_func(key, operator.getitem) |
| 62 | and query.args[0] == key.args[0] |
| 63 | and is_func(query.args[0], torch.ops.aten.split_with_sizes.default) |
| 64 | and all( |
| 65 | is_func(user, torch.ops.aten.slice_scatter.default) |
| 66 | for getitem_node in getitem_nodes.values() |
| 67 | for user in getitem_node.users |
| 68 | ) |
| 69 | ): |
| 70 | # Pattern where query and key are slices of an mm_node. |
| 71 | # While functionalized, results at [1] and [2] are scattered |
| 72 | # back into mm_node. So after de-functionalization, we can |
| 73 | # just use mm_node directly. |
| 74 | |
| 75 | mm_node = query.args[0].args[0] |
| 76 | for user in getitem_nodes.values(): |
| 77 | for user_of_getitem in user.users: |
| 78 | if is_func( |
| 79 | user_of_getitem, torch.ops.aten.slice_scatter.default |
| 80 | ): |
| 81 | user_of_getitem.replace_all_uses_with(mm_node) |
| 82 | self._remove(user_of_getitem) |
| 83 | self._remove(user) |
| 84 | |
| 85 | self.insert_defunctionalized(graph, node) |
| 86 | self._remove(node) |
| 87 | |
| 88 | else: |
| 89 | # Directly replace the auto_functionalize(rotary_embedding) |
| 90 | # with the inplace rotary_embedding. In theory, we shouldn't |
| 91 | # do this blindly, but in practice in vLLM it's ok. The best |
| 92 | # solution is to use auto_functionalization_v2 and then use |
| 93 | # inductor's builtin defunctionalization (reinplacing) pass. |
| 94 | mutated_args = {1: "query", 2: "key"} |
| 95 | self.defunctionalize(graph, node, mutated_args) |
| 96 | |
| 97 | # rms_norm replacements avoid the most copies for LLaMa. |
| 98 | elif at_target == torch.ops._C.fused_add_rms_norm.default: |
| 99 | mutated_args = {1: "input", 2: "residual"} |
| 100 | self.defunctionalize(graph, node, mutated_args) |
| 101 | elif at_target == torch.ops._C.fused_add_rms_norm_static_fp8_quant.default: # noqa: E501 |
| 102 | mutated_args = {1: "result", 2: "residual"} |
| 103 | self.defunctionalize(graph, node, mutated_args) |
| 104 | elif at_target == torch.ops._C.rms_norm_dynamic_per_token_quant.default: # noqa: E501 |
| 105 | mutated_args = {1: "result", 2: "scale", 3: "residual"} |
| 106 | self.defunctionalize(graph, node, mutated_args) |
| 107 | elif at_target in [ |
| 108 | torch.ops._C.rms_norm.default, |
| 109 | torch.ops._C.rms_norm_static_fp8_quant.default, |
| 110 | ]: |
| 111 | mutated_args = {1: "result"} |
| 112 | self.defunctionalize(graph, node, mutated_args) |
| 113 | elif ( |
| 114 | hasattr(torch.ops.vllm, "flashinfer_trtllm_fused_allreduce_norm") |
| 115 | and at_target |
| 116 | == torch.ops.vllm.flashinfer_trtllm_fused_allreduce_norm.default |
| 117 | ): |
| 118 | mutated_args = { |
| 119 | 1: "allreduce_in", |
| 120 | 2: "residual", |
| 121 | 3: "norm_out", |
| 122 | 4: "quant_out", |
| 123 | 5: "scale_out", |
| 124 | } |
| 125 | self.defunctionalize(graph, node, mutated_args) |
| 126 | elif ( |
| 127 | hasattr( |
| 128 | torch.ops.vllm, |
| 129 | "flashinfer_fused_add_rms_norm_nvfp4_quant", |
| 130 | ) |
| 131 | and at_target |
| 132 | == torch.ops.vllm.flashinfer_fused_add_rms_norm_nvfp4_quant.default |
| 133 | ): |
| 134 | mutated_args = { |
| 135 | 1: "result", |
| 136 | 2: "result_block_scale", |
| 137 | 3: "residual", |
| 138 | } |
| 139 | self.defunctionalize(graph, node, mutated_args) |
| 140 | # For some reason we need to specify the args for both |
| 141 | # silu_and_mul and silu_and_mul_quant. The kwargs |
| 142 | # pathway gets the wrong answer. |
| 143 | elif at_target == torch.ops._C.silu_and_mul.default: |
| 144 | mutated_args = {1: "result"} |
| 145 | self.defunctionalize( |
| 146 | graph, node, mutated_args, args=("result", "input") |
| 147 | ) |
| 148 | elif at_target == torch.ops._C.silu_and_mul_quant.default: |
| 149 | mutated_args = {1: "result"} |
| 150 | self.defunctionalize( |
| 151 | graph, node, mutated_args, args=("result", "input", "scale") |
| 152 | ) |
| 153 | elif ( |
| 154 | hasattr(torch.ops._C, "silu_and_mul_nvfp4_quant") |
| 155 | and at_target == torch.ops._C.silu_and_mul_nvfp4_quant.default |
| 156 | ): |
| 157 | mutated_args = {1: "result", 2: "result_block_scale"} |
| 158 | self.defunctionalize( |
| 159 | graph, |
| 160 | node, |
| 161 | mutated_args, |
| 162 | args=( |
| 163 | "result", |
| 164 | "result_block_scale", |
| 165 | "input", |
| 166 | "input_global_scale", |
| 167 | ), |
| 168 | ) |
| 169 | # Defunctionalize fused_qk_norm_rope to remove higher-order wrapper. |
| 170 | elif at_target == torch.ops._C.fused_qk_norm_rope.default: |
| 171 | mutated_args = {1: "qkv"} |
| 172 | args = ( |
| 173 | "qkv", |
| 174 | "num_heads_q", |
| 175 | "num_heads_k", |
| 176 | "num_heads_v", |
| 177 | "head_dim", |
| 178 | "eps", |
| 179 | "q_weight", |
| 180 | "k_weight", |
| 181 | "cos_sin_cache", |
| 182 | "is_neox", |
| 183 | "position_ids", |
| 184 | "forced_token_heads_per_warp", |
| 185 | ) |
| 186 | self.defunctionalize(graph, node, mutated_args=mutated_args, args=args) |
| 187 | elif ( |
| 188 | hasattr(torch.ops.vllm, "fused_rope_and_unified_kv_cache_update") |
| 189 | and at_target |
| 190 | == torch.ops.vllm.fused_rope_and_unified_kv_cache_update.default |
| 191 | ): |
| 192 | mutated_args = { |
| 193 | 1: "query", |
| 194 | 2: "key", |
| 195 | } |
| 196 | self.defunctionalize(graph, node, mutated_args=mutated_args) |
| 197 | elif ( |
| 198 | hasattr(torch.ops.vllm, "fused_rope_unified_mla_kv_cache_update") |
| 199 | and at_target |
| 200 | == torch.ops.vllm.fused_rope_unified_mla_kv_cache_update.default |
| 201 | ): |
| 202 | # AOTAutograd functionalizes `q[..., nope_dim:] = rope_result` into |
| 203 | # a sequence of aten ops on q: view+slice+copy+slice_scatter. |
| 204 | # Since the fused MLA RoPE op mutates q_pe in-place, we can remove |
| 205 | # the redundant copy and slice_scatter ops during defunctionalization. |
| 206 | getitem_nodes = self.getitem_users(node) |
| 207 | q_pe_out = getitem_nodes[1] |
| 208 | |
| 209 | for user in list(q_pe_out.users): |
| 210 | if is_func(user, torch.ops.aten.copy.default): |
| 211 | copy_temp = user |
| 212 | slice_temp = copy_temp.args[0] |
| 213 | for user in list(copy_temp.users): |
| 214 | if is_func(user, torch.ops.aten.slice_scatter.default): |
| 215 | slice_scatter_temp = user |
| 216 | view_temp = slice_scatter_temp.args[0] |
| 217 | |
| 218 | view_orig = slice_temp.args[0] |
| 219 | slice_scatter_temp.replace_all_uses_with(view_orig) |
| 220 | self._remove(slice_scatter_temp) |
| 221 | self._remove(copy_temp) |
| 222 | self._remove(slice_temp) |
| 223 | self._remove(view_temp) |
| 224 | self._remove(q_pe_out) |
| 225 | |
| 226 | # defunctionalize k_pe manually; self.replace_users_with_mutated_args |
| 227 | # does not support only replacing specific kwargs |
| 228 | k_pe_in = node.kwargs["k_pe"] |
| 229 | k_pe_out = getitem_nodes[2] |
| 230 | k_pe_out.replace_all_uses_with(k_pe_in) |
| 231 | self._remove(k_pe_out) |
| 232 | |
| 233 | self.insert_defunctionalized(graph, node) |
| 234 | self._remove(node) |
| 235 | |
| 236 | # only used for test_functionalization::TestFunctionWithMutatedArgsAndReturn |
| 237 | elif ( |
| 238 | hasattr(torch.ops.vllm, "function_with_mutated_args_and_return") |
| 239 | and at_target |
| 240 | == torch.ops.vllm.function_with_mutated_args_and_return.default |
| 241 | ): |
| 242 | mutated_args = {1: "x"} |
| 243 | self.defunctionalize(graph, node, mutated_args=mutated_args) |
| 244 | else: |
| 245 | continue # skip the count |
| 246 | |
| 247 | count += 1 |
| 248 | |
| 249 | self.dump_graph(graph, "before_cleanup") |
| 250 | |
| 251 | # Remove the nodes all at once |
| 252 | count_removed = len(self.nodes_to_remove) |
| 253 | for node in self.nodes_to_remove: |
| 254 | graph.erase_node(node) |
| 255 | |
| 256 | logger.debug( |
| 257 | "De-functionalized %s nodes, removed %s nodes", count, count_removed |
| 258 | ) |
| 259 | self.nodes_to_remove.clear() |
| 260 | |
| 261 | def _remove(self, node_or_nodes: torch.fx.Node | Iterable[torch.fx.Node]) -> None: |
| 262 | """Stage a node (or nodes) for removal at the end of the pass.""" |
| 263 | if isinstance(node_or_nodes, torch.fx.Node): |
| 264 | self.nodes_to_remove.append(node_or_nodes) |
| 265 | else: |
| 266 | self.nodes_to_remove.extend(node_or_nodes) |
| 267 | |
| 268 | def defunctionalize( |
| 269 | self, |
| 270 | graph: torch.fx.Graph, |
| 271 | node: torch.fx.Node, |
| 272 | mutated_args: dict[int, torch.fx.Node | str], |
| 273 | args: tuple[torch.fx.Node | str, ...] | None = None, |
| 274 | ) -> None: |
| 275 | """De-functionalize a node by replacing it with a call to the original. |
| 276 | It also replaces the getitem users with the mutated arguments. |
| 277 | See replace_users_with_mutated_args and insert_defunctionalized. |
| 278 | """ |
| 279 | self.replace_users_with_mutated_args(node, mutated_args) |
| 280 | self.insert_defunctionalized(graph, node, args=args) |
| 281 | self._remove(node) |
| 282 | |
| 283 | def replace_users_with_mutated_args( |
| 284 | self, node: torch.fx.Node, mutated_args: dict[int, torch.fx.Node | str] |
| 285 | ) -> None: |
| 286 | """Replace mutated getitem users of the auto-functionalized node with the |
| 287 | mutated arguments. |
| 288 | |
| 289 | Args: |
| 290 | node: The auto-functionalized node |
| 291 | mutated_args: The mutated arguments, indexed by getitem index. |
| 292 | If the value of an arg is a string, `node.kwargs[arg]` is used. |
| 293 | |
| 294 | """ |
| 295 | for idx, user in self.getitem_users(node).items(): |
| 296 | # Some functionalized nodes may return both a result at getitem[0] |
| 297 | # as well as mutated args at getitem[1:...] |
| 298 | if idx == 0: |
| 299 | assert idx not in mutated_args, ( |
| 300 | f"result at getitem[0] should not be in mutated_args for {node}" |
| 301 | ) |
| 302 | continue |
| 303 | arg = mutated_args[idx] |
| 304 | arg = node.kwargs[arg] if isinstance(arg, str) else arg |
| 305 | user.replace_all_uses_with(arg) |
| 306 | self._remove(user) |
| 307 | |
| 308 | def getitem_users(self, node: torch.fx.Node) -> dict[int, torch.fx.Node]: |
| 309 | """Returns the operator.getitem users of the auto-functionalized node, |
| 310 | indexed by the index they are getting. |
| 311 | """ |
| 312 | users = {} |
| 313 | for user in node.users: |
| 314 | if is_func(user, operator.getitem): |
| 315 | idx = user.args[1] |
| 316 | users[idx] = user |
| 317 | return users |
| 318 | |
| 319 | def insert_defunctionalized( |
| 320 | self, |
| 321 | graph: torch.fx.Graph, |
| 322 | node: torch.fx.Node, |
| 323 | args: tuple[torch.fx.Node | str, ...] | None = None, |
| 324 | ) -> None: |
| 325 | """Insert a new defunctionalized node into the graph before node. |
| 326 | If one of the kwargs is 'out', provide args directly, |
| 327 | as node.kwargs cannot be used. |
| 328 | See https://github.com/pytorch/pytorch/blob/a00faf440888ffb724bad413f329a49e2b6388e7/torch/_inductor/lowering.py#L351 |
| 329 | |
| 330 | Args: |
| 331 | graph: Graph to insert the defunctionalized node into |
| 332 | node: The auto-functionalized node to defunctionalize |
| 333 | args: If we cannot use kwargs, specify args directly. |
| 334 | If an arg is a string, `node.kwargs[arg]` is used. |
| 335 | |
| 336 | """ # noqa: E501 |
| 337 | assert is_func(node, auto_functionalized), ( |
| 338 | f"node must be auto-functionalized, is {node} instead" |
| 339 | ) |
| 340 | |
| 341 | # Create a new call to the original function |
| 342 | with graph.inserting_before(node): |
| 343 | function = node.args[0] |
| 344 | if args is None: |
| 345 | fn_node = graph.call_function(function, kwargs=node.kwargs) |
| 346 | else: |
| 347 | # Args passed as strings refer to items in node.kwargs |
| 348 | args = tuple( |
| 349 | node.kwargs[arg] if isinstance(arg, str) else arg for arg in args |
| 350 | ) |
| 351 | fn_node = graph.call_function(function, args=args) |
| 352 | |
| 353 | # If the function returns a value as well as mutating args inplace, |
| 354 | # the functionalized node will have a getitem[0] user that holds this value |
| 355 | # Replace getitem[0] user of the auto-functionalized node |
| 356 | # with the new defunctionalized node directly if it exists |
| 357 | users = self.getitem_users(node) |
| 358 | if 0 in users: |
| 359 | user = users[0] |
| 360 | user.replace_all_uses_with(fn_node) |
| 361 | self._remove(user) |