| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections import defaultdict |
| 4 | |
| 5 | from torch import fx |
| 6 | from torch._inductor.pattern_matcher import ( |
| 7 | PatternMatcherPass, |
| 8 | ) |
| 9 | |
| 10 | from vllm.config import VllmConfig |
| 11 | from vllm.logger import init_logger |
| 12 | |
| 13 | from ..inductor_pass import get_pass_context |
| 14 | from ..vllm_inductor_pass import VllmInductorPass |
| 15 | from .lowering_pass import get_ir_op |
| 16 | from .utils import overload_or_default |
| 17 | |
| 18 | logger = init_logger(__name__) |
| 19 | |
| 20 | |
| 21 | class VllmIRInplaceFunctionalizationPass(VllmInductorPass): |
| 22 | """This pass functionalizes maybe_inplace vLLM IR ops to the default overload. |
| 23 | The maybe_inplace overloads have the same signature as the default overload |
| 24 | so the pass simply replaces the called overload. |
| 25 | That makes the graph properly functional. |
| 26 | The pass also validates that activations passed to maybe_inplace have no later |
| 27 | uses in the graph: they are donated to the maybe_inplace op call, |
| 28 | and their contents are not defined afterward. |
| 29 | |
| 30 | This pass operates pre-AOTAutograd, |
| 31 | so it must handle non-normalized and non-functional IR. |
| 32 | """ |
| 33 | |
| 34 | def __init__(self, vllm_config: VllmConfig) -> None: |
| 35 | super().__init__(vllm_config) |
| 36 | self.patterns = PatternMatcherPass(self.pass_name) |
| 37 | self.functionalized_ops: dict[str, int] = defaultdict(lambda: 0) |
| 38 | |
| 39 | @VllmInductorPass.time_and_log |
| 40 | def __call__(self, graph: fx.Graph) -> None: |
| 41 | # clear at the beginning instead of end, so that tests can inspect |
| 42 | self.functionalized_ops.clear() |
| 43 | assert graph.owning_module is not None |
| 44 | node_to_idx = {node: i for i, node in enumerate(graph.nodes)} |
| 45 | |
| 46 | # Pass donated input via vLLM's pass context |
| 47 | pass_context = get_pass_context() |
| 48 | pass_context.donated_input_ids = set[int]() |
| 49 | |
| 50 | for node in graph.nodes: |
| 51 | if (ir_op := get_ir_op(node)) is None: |
| 52 | continue |
| 53 | |
| 54 | op_overload = overload_or_default(node.target) |
| 55 | overload_name = op_overload._overloadname |
| 56 | if overload_name != "maybe_inplace": |
| 57 | assert overload_name == "default", ( |
| 58 | f"Found overload {overload_name} for op {ir_op.name}, " |
| 59 | f"expected maybe_inplace or default" |
| 60 | ) |
| 61 | continue |
| 62 | |
| 63 | # must have maybe_inplace overload and allow_inplace |
| 64 | assert ir_op.allow_inplace and hasattr(ir_op, "maybe_inplace") |
| 65 | |
| 66 | # Check that activation inputs are not used after this op |
| 67 | for arg_idx in ir_op.activation_indices: |
| 68 | arg = node.args[arg_idx] |
| 69 | assert isinstance(arg, fx.Node), "Activation inputs must be fx.Node" |
| 70 | for user in arg.users: |
| 71 | if node_to_idx[user] > node_to_idx[node]: |
| 72 | raise ValueError( |
| 73 | f"Input {arg} to maybe_inplace node {node} " |
| 74 | f"is used again after the node. " |
| 75 | f"This is not allowed; activation inputs to maybe_inplace " |
| 76 | f"ops are donated to the op, meaning their memory may be " |
| 77 | f"recycled for outputs.\n\n" |
| 78 | f"To preserve the inputs, use the default overload or " |
| 79 | f"clone them manually beforehand." |
| 80 | ) |
| 81 | |
| 82 | if arg.op == "placeholder": |
| 83 | # Graph input that maybe_inplace might modify. |
| 84 | # Mark it so downstream passes know it's donated. |
| 85 | # TODO(luka) store in placeholder node meta once supported |
| 86 | pass_context.donated_input_ids.add(node_to_idx[arg]) |
| 87 | |
| 88 | # Same signature, just replace the overload that's called. |
| 89 | node.target = ir_op.torch_op |
| 90 | self.functionalized_ops[ir_op.name] += 1 |
| 91 | |
| 92 | count = sum(self.functionalized_ops.values()) |
| 93 | ops = ",".join(self.functionalized_ops.keys()) |
| 94 | logger.debug("Donated input IDs: %s", pass_context.donated_input_ids) |
| 95 | logger.debug( |
| 96 | "%s functionalized %d vLLM IR nodes for op(s) %s", |
| 97 | self.pass_name, |
| 98 | count, |
| 99 | ops, |
| 100 | ) |