| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | from collections import defaultdict |
| 4 | from collections.abc import Iterable |
| 5 | |
| 6 | from torch import fx |
| 7 | from torch._inductor.pattern_matcher import ( |
| 8 | CallFunctionVarArgs, |
| 9 | Match, |
| 10 | PatternMatcherPass, |
| 11 | register_graph_pattern, |
| 12 | ) |
| 13 | |
| 14 | from vllm.config import VllmConfig |
| 15 | from vllm.ir.op import IrOp |
| 16 | from vllm.logger import init_logger |
| 17 | from vllm.logging_utils import lazy |
| 18 | |
| 19 | from ..vllm_inductor_pass import VllmInductorPass |
| 20 | from .utils import get_ir_op |
| 21 | |
| 22 | logger = init_logger(__name__) |
| 23 | |
| 24 | |
| 25 | class VllmIRLoweringPass(VllmInductorPass): |
| 26 | """This pass lowers vLLM IR ops to their implementations the priority list.""" |
| 27 | |
| 28 | def __init__(self, vllm_config: VllmConfig) -> None: |
| 29 | super().__init__(vllm_config) |
| 30 | self.patterns = PatternMatcherPass(self.pass_name) |
| 31 | self.selected_impls: dict[str, dict[str, str]] = defaultdict(lambda: {}) |
| 32 | self.ops = [ir_op.torch_op for ir_op in IrOp.registry.values()] |
| 33 | |
| 34 | # Look for any call_function node where the target is a vLLM IR op. |
| 35 | # Then, lower_matched_op will select, trace, and insert the implementation. |
| 36 | register_graph_pattern( |
| 37 | CallFunctionVarArgs(self.ops), |
| 38 | pass_dict=self.patterns, |
| 39 | )(self.lower_matched_op) |
| 40 | |
| 41 | def lower_matched_op(self, match: Match, *args, **kwargs): |
| 42 | # TODO(luka) I think args and kwargs are for the match, but just use the node? |
| 43 | |
| 44 | assert len(match.nodes) == 1, "Expected single node match" |
| 45 | node = match.nodes[0] |
| 46 | ir_op = get_ir_op(node) # TODO is node.target always an overload? |
| 47 | assert ir_op is not None, "Expected vLLM IR op" |
| 48 | assert not node.kwargs # I think there should never be kwargs here |
| 49 | |
| 50 | # Select and record the implementation, using fake args |
| 51 | fake_args = fx.map_arg(node.args, lambda arg: arg.meta["val"]) |
| 52 | ir_op_impl = ir_op.dispatch(*fake_args) |
| 53 | self.selected_impls[ir_op.name][node.name] = ir_op_impl.provider |
| 54 | |
| 55 | # replace_by_example wants node args, not the fake tensors |
| 56 | # use func_impl_fn to properly handle in-place implementations |
| 57 | # TODO(luka): Use aot_export_module to get functionalized graph |
| 58 | # TODO(luka): Cache the fx_replacement to avoid re-tracing the same impl |
| 59 | |
| 60 | # Defaults not present on node.args but required for replacement tracing |
| 61 | bound_args = ir_op._py_signature.bind(*node.args) |
| 62 | bound_args.apply_defaults() |
| 63 | # It is not safe to run functional passes (like DCE) on the replacements |
| 64 | # as they might not be functional. |
| 65 | match.replace_by_example( |
| 66 | ir_op_impl.func_impl_fn, bound_args.args, run_functional_passes=False |
| 67 | ) |
| 68 | |
| 69 | @VllmInductorPass.time_and_log |
| 70 | def __call__(self, graph: fx.Graph) -> None: |
| 71 | # clear at the beginning instead of end, so that tests can inspect |
| 72 | self.selected_impls.clear() |
| 73 | |
| 74 | count = self.patterns.apply(graph) |
| 75 | logger.debug("VllmIRLoweringPass lowered %d vLLM IR nodes", count) |
| 76 | |
| 77 | # TODO write self.selected_impls to depyf/tlparse dir |
| 78 | def count_items(impls: Iterable[str]) -> dict[str, int]: |
| 79 | counts: dict[str, int] = defaultdict(lambda: 0) |
| 80 | for impl in impls: |
| 81 | counts[impl] += 1 |
| 82 | return counts |
| 83 | |
| 84 | def print_count(counts: dict[str, int]) -> str: |
| 85 | # e.g., "impl1*3,impl2" |
| 86 | impl_count = lambda i, c: f"{i}" if c == 1 else f"{i}*{c}" |
| 87 | return ",".join(impl_count(impl, count) for impl, count in counts.items()) |
| 88 | |
| 89 | logger.debug( |
| 90 | "Selected implementations: %s", |
| 91 | lazy( |
| 92 | lambda: ", ".join( |
| 93 | f"{op}={print_count(count_items(impls_by_node.values()))}" |
| 94 | for op, impls_by_node in self.selected_impls.items() |
| 95 | ) |
| 96 | ), |
| 97 | ) |
| 98 | |
| 99 | failed_nodes: list[fx.Node] = [] |
| 100 | failed_ops: set[str] = set() |
| 101 | # Check no vllm_ir nodes were left in the graph |
| 102 | for node in graph.nodes: |
| 103 | if (ir_op := get_ir_op(node)) is None: |
| 104 | continue |
| 105 | |
| 106 | failed_nodes.append(node) |
| 107 | failed_ops.add(ir_op.name) |
| 108 | |
| 109 | if failed_nodes or failed_ops: |
| 110 | logger.warning("Failed to lower vLLM IR ops: %s", ",".join(failed_ops)) |
| 111 | logger.warning("Full node list: %s", ",".join(str(n) for n in failed_nodes)) |
| 112 | |
| 113 | def uuid(self) -> str: |
| 114 | """IR op priority & impl sources affect lowering pass output, |
| 115 | so we include them in the cache key. |
| 116 | """ |
| 117 | priorities = {name: op.get_priority() for name, op in IrOp.registry.items()} |
| 118 | priorities_str = ";".join( |
| 119 | f"{name}={','.join(p)}" for name, p in priorities.items() |
| 120 | ) |
| 121 | |
| 122 | impl_uuids_str = ";".join( |
| 123 | f"{name}=" |
| 124 | + ",".join(IrOp.registry[name].impls[provider].uuid() for provider in p) |
| 125 | for name, p in priorities.items() |
| 126 | ) |
| 127 | |
| 128 | return f"{super().uuid()}|{priorities_str}|{impl_uuids_str}" |