| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import weakref |
| 5 | from collections.abc import Callable, Sequence |
| 6 | from contextlib import nullcontext |
| 7 | from copy import deepcopy |
| 8 | |
| 9 | import depyf |
| 10 | from torch import fx |
| 11 | from torch._ops import OpOverload, OpOverloadPacket |
| 12 | from torch.fx._utils import lazy_format_graph_code |
| 13 | |
| 14 | from vllm.compilation.passes.fx_utils import find_op_nodes |
| 15 | from vllm.compilation.passes.inductor_pass import ( |
| 16 | InductorPass, |
| 17 | pass_context, |
| 18 | ) |
| 19 | from vllm.compilation.passes.ir.inplace_functionalization import ( |
| 20 | VllmIRInplaceFunctionalizationPass, |
| 21 | ) |
| 22 | from vllm.compilation.passes.pass_manager import with_pattern_match_debug |
| 23 | from vllm.compilation.passes.vllm_inductor_pass import VllmInductorPass |
| 24 | from vllm.config import VllmConfig, get_current_vllm_config |
| 25 | from vllm.config.utils import Range |
| 26 | from vllm.logger import init_logger |
| 27 | |
| 28 | logger = init_logger("vllm.tests.compile.backend") |
| 29 | |
| 30 | |
| 31 | class LazyInitPass(InductorPass): |
| 32 | """If there's a pass that we want to initialize lazily in a test, |
| 33 | we can wrap it in LazyInitPass, which will initialize the pass when invoked |
| 34 | and then immediately invoke it. |
| 35 | """ |
| 36 | |
| 37 | def __init__(self, pass_cls: type[VllmInductorPass], vllm_config: VllmConfig): |
| 38 | self.pass_cls = pass_cls |
| 39 | self.vllm_config = weakref.proxy(vllm_config) # avoid cycle |
| 40 | |
| 41 | def __call__(self, graph: fx.Graph) -> None: |
| 42 | self.pass_ = self.pass_cls(self.vllm_config) |
| 43 | self.pass_(graph) |
| 44 | |
| 45 | |
| 46 | class TestBackend: |
| 47 | """This class provides a simple Inductor backend that can be used for testing. |
| 48 | It takes a list of custom passes and runs them after Inductor's passes. |
| 49 | It also saves the graph before and after the custom passes for inspection. |
| 50 | |
| 51 | Inductor config can be modified directly by editing the inductor_config |
| 52 | property. This can be helpful for adding passes like the |
| 53 | 'pre_grad_custom_pass' and the 'post_grad_custom_pre_pass'. |
| 54 | Inductor config is default-initialized from VllmConfig.CompilationConfig. |
| 55 | """ |
| 56 | |
| 57 | def __init__(self, *passes: InductorPass | Callable[[fx.Graph], None]): |
| 58 | self.custom_passes = list(passes) |
| 59 | vllm_config = get_current_vllm_config() |
| 60 | compile_config = vllm_config.compilation_config |
| 61 | self.range = Range(1, vllm_config.scheduler_config.max_num_batched_tokens) |
| 62 | # Deepcopy to allow multiple TestBackend instances to use the same VllmConfig |
| 63 | self.inductor_config = deepcopy(compile_config.inductor_compile_config) |
| 64 | self.inductor_config["force_disable_caches"] = True |
| 65 | self.inductor_config["post_grad_custom_post_pass"] = self.post_pass |
| 66 | |
| 67 | # Add VllmIRInplaceFunctionalizationPass as pre-grad pass by default |
| 68 | self.inductor_config["pre_grad_custom_pass"] = ( |
| 69 | VllmIRInplaceFunctionalizationPass(vllm_config) |
| 70 | ) |
| 71 | |
| 72 | if debug_dump_path := vllm_config.compile_debug_dump_path(): |
| 73 | logger.debug("Dumping depyf output to %s", debug_dump_path) |
| 74 | self.debug_ctx = depyf.prepare_debug(debug_dump_path.as_posix()) |
| 75 | else: |
| 76 | self.debug_ctx = nullcontext() |
| 77 | |
| 78 | def __call__(self, graph: fx.GraphModule, example_inputs): |
| 79 | self.graph_pre_compile = deepcopy(graph) |
| 80 | from torch._inductor.compile_fx import compile_fx |
| 81 | |
| 82 | with self.debug_ctx, pass_context(self.range): |
| 83 | return compile_fx( |
| 84 | graph, example_inputs, config_patches=self.inductor_config |
| 85 | ) |
| 86 | |
| 87 | @with_pattern_match_debug |
| 88 | def post_pass(self, graph: fx.Graph): |
| 89 | self.graph_pre_pass = deepcopy(graph) |
| 90 | lazy_format_graph_code("graph_pre_pass", graph.owning_module) |
| 91 | |
| 92 | VllmInductorPass.dump_prefix = 0 |
| 93 | for pass_ in self.custom_passes: |
| 94 | pass_(graph) |
| 95 | VllmInductorPass.dump_prefix += 1 |
| 96 | |
| 97 | VllmInductorPass.dump_prefix = None |
| 98 | |
| 99 | self.graph_post_pass = deepcopy(graph) |
| 100 | lazy_format_graph_code("graph_post_pass", graph.owning_module) |
| 101 | # assign by reference, will reflect the final state of the graph |
| 102 | self.final_graph = graph |
| 103 | |
| 104 | def check_before_ops( |
| 105 | self, ops: Sequence[OpOverload | OpOverloadPacket], fully_replaced=True |
| 106 | ): |
| 107 | for op in ops: |
| 108 | num_pre = len(list(find_op_nodes(op, self.graph_pre_pass))) |
| 109 | num_post = len(list(find_op_nodes(op, self.graph_post_pass))) |
| 110 | assert num_pre > 0, f"Op {op.name()} not found in pre-pass graph" |
| 111 | assert num_pre > num_post, f"All nodes remain for op {op.name()}" |
| 112 | if fully_replaced: |
| 113 | assert num_post == 0, f"Unexpected op {op.name()} in post-pass graph" |
| 114 | |
| 115 | def check_after_ops(self, ops: Sequence[OpOverload | OpOverloadPacket]): |
| 116 | for op in ops: |
| 117 | num_pre = len(list(find_op_nodes(op, self.graph_pre_pass))) |
| 118 | num_post = len(list(find_op_nodes(op, self.graph_post_pass))) |
| 119 | assert num_pre == 0, f"Unexpected op {op.name()} in pre-pass graph" |
| 120 | assert num_post > 0, f"Op {op.name()} not found in post-pass graph" |
| 121 | |
| 122 | def op_count(self, op: OpOverload | OpOverloadPacket, before=False) -> int: |
| 123 | graph = self.graph_pre_pass if before else self.graph_post_pass |
| 124 | return len(list(find_op_nodes(op, graph))) |
| 125 | |
| 126 | def print_graphs(self): |
| 127 | print("=== Graph before custom passes ===") |
| 128 | print(self.graph_pre_pass.python_code(root_module="self", verbose=True).src) |
| 129 | print("=== Graph after custom passes ===") |
| 130 | print(self.graph_post_pass.python_code(root_module="self", verbose=True).src) |