| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | import functools |
| 4 | import operator |
| 5 | import time |
| 6 | from abc import ABC, abstractmethod |
| 7 | from collections import defaultdict |
| 8 | from collections.abc import Callable |
| 9 | from dataclasses import dataclass |
| 10 | from typing import Any, ClassVar, Generic, ParamSpec, TypeVar |
| 11 | |
| 12 | import regex as re |
| 13 | import torch |
| 14 | import torch._inductor.pattern_matcher as pm |
| 15 | from torch import fx |
| 16 | from torch._dynamo.utils import lazy_format_graph_code |
| 17 | from torch._inductor.pattern_matcher import PatternMatcherPass, PatternPrettyPrinter |
| 18 | |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.logger import init_logger |
| 21 | from vllm.platforms import current_platform |
| 22 | |
| 23 | from .fx_utils import is_func |
| 24 | from .inductor_pass import InductorPass, enable_fake_mode |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | DEVICE_TYPE = current_platform.device_type |
| 29 | |
| 30 | |
| 31 | @dataclass |
| 32 | class InductorCompilationConfig: |
| 33 | splitting_ops: list[str] | None = None |
| 34 | use_inductor_graph_partition: bool = False |
| 35 | |
| 36 | |
| 37 | class VllmInductorPass(InductorPass): |
| 38 | """An inductor pass with access to vLLM PassConfig. |
| 39 | It provides timing, logging, and dumping utilities. |
| 40 | """ |
| 41 | |
| 42 | dump_prefix: ClassVar[int | None] = None |
| 43 | """Keep track of pass index for debug dump ordering.""" |
| 44 | |
| 45 | def __init__(self, config: VllmConfig): |
| 46 | # Get only the necessary CompilationConfig for the inductor pass, since |
| 47 | # full `CompilationConfig` contains pointer to model which is unsafe. |
| 48 | self.compilation_config = InductorCompilationConfig( |
| 49 | splitting_ops=config.compilation_config.splitting_ops, |
| 50 | use_inductor_graph_partition=config.compilation_config.use_inductor_graph_partition, |
| 51 | ) |
| 52 | self.pass_config = config.compilation_config.pass_config |
| 53 | self.model_dtype = config.model_config.dtype if config.model_config else None |
| 54 | self.device: str | None = ( |
| 55 | config.device_config.device if config.device_config else None |
| 56 | ) |
| 57 | self.pass_name = self.__class__.__name__ |
| 58 | |
| 59 | @staticmethod |
| 60 | def time_and_log( |
| 61 | call_fn: Callable[["VllmInductorPass", torch.fx.Graph], None], |
| 62 | ) -> Callable[["VllmInductorPass", torch.fx.Graph], None]: |
| 63 | @functools.wraps(call_fn) |
| 64 | def wrapped(self: VllmInductorPass, graph: torch.fx.Graph) -> None: |
| 65 | self.begin() |
| 66 | self.dump_graph(graph, "before") |
| 67 | call_fn(self, graph) |
| 68 | self.dump_graph(graph, "after") |
| 69 | self.end_and_log() |
| 70 | |
| 71 | return wrapped |
| 72 | |
| 73 | def dump_graph(self, graph: torch.fx.Graph, stage: str) -> None: |
| 74 | i = VllmInductorPass.dump_prefix |
| 75 | i_str = "" if i is None else f".{i}" |
| 76 | lazy_format_graph_code( |
| 77 | f"post_grad{i_str}.{self.pass_name}.{stage}", graph.owning_module |
| 78 | ) |
| 79 | |
| 80 | def begin(self) -> None: |
| 81 | self._start_time = time.perf_counter_ns() |
| 82 | |
| 83 | def end_and_log(self) -> None: |
| 84 | self._end_time = time.perf_counter_ns() |
| 85 | duration_ms = float(self._end_time - self._start_time) / 1.0e6 |
| 86 | logger.debug("%s completed in %.1f ms", self.pass_name, duration_ms) |
| 87 | |
| 88 | |
| 89 | def get_match_table() -> dict[str, int]: |
| 90 | """Return a snapshot of the match table.""" |
| 91 | return dict(VllmPatternMatcherPass.match_table) |
| 92 | |
| 93 | |
| 94 | class VllmPatternMatcherPass(VllmInductorPass): |
| 95 | """A VllmInductorPass that uses the Inductor pattern matcher. |
| 96 | Provides pattern registration with match counting, debug dumping, and logging. |
| 97 | """ |
| 98 | |
| 99 | matched_count: int = 0 |
| 100 | """The number of matched patterns in the pass.""" |
| 101 | |
| 102 | match_table: ClassVar[defaultdict[str, int]] = defaultdict(int) |
| 103 | """Global table mapping pass name to its total match count.""" |
| 104 | |
| 105 | _OP_OVERLOAD_PATTERN: ClassVar[re.Pattern] = re.compile( |
| 106 | r"<OpOverload\(op='([^']*)', overload='([^']*)'\)>" |
| 107 | ) |
| 108 | |
| 109 | def _replace_op_overloads(self, string: str) -> str: |
| 110 | """Replace <OpOverload(..., ...)> with nicer formulations.""" |
| 111 | return str( |
| 112 | self._OP_OVERLOAD_PATTERN.sub( |
| 113 | lambda m: f"torch.ops.{m.group(1)}.{m.group(2)}", |
| 114 | string, |
| 115 | ) |
| 116 | ) |
| 117 | |
| 118 | @classmethod |
| 119 | def log_match_summary(cls) -> None: |
| 120 | if cls.match_table: |
| 121 | logger.debug("fusion pass matches: %s", dict(cls.match_table)) |
| 122 | |
| 123 | def dump_patterns(self, config: VllmConfig, pm_pass: PatternMatcherPass) -> None: |
| 124 | """If debug dumping is enabled, dump the Inductor pattern-matcher patterns |
| 125 | into the debug_dump_path folder next to the dumped fx graphs. |
| 126 | |
| 127 | This method does its best to print something that looks like Python code |
| 128 | for easier debugging and potentially navigation. If any errors appear in |
| 129 | the output, please add to this method. |
| 130 | |
| 131 | TODO(luka): use pattern object to manually produce pattern graph |
| 132 | """ |
| 133 | debug_dump_path = config.compile_debug_dump_path() |
| 134 | if not debug_dump_path: |
| 135 | return |
| 136 | |
| 137 | debug_dump_path.mkdir(parents=True, exist_ok=True) |
| 138 | |
| 139 | from vllm.utils.system_utils import unique_filepath |
| 140 | |
| 141 | file_path = unique_filepath( |
| 142 | lambda i: debug_dump_path / f"patterns.{self.pass_name}.{i}.py" |
| 143 | ) |
| 144 | |
| 145 | with file_path.open("w") as f: |
| 146 | print( |
| 147 | f"# This file was produced by VllmPatternMatcherPass." |
| 148 | f"dump_patterns for {self.pass_name}.\n" |
| 149 | f"# It does its best to produce valid-Python-looking code but" |
| 150 | f" please add to dump_patterns if there are any errors.\n\n" |
| 151 | f"from torch._higher_order_ops.auto_functionalize import " |
| 152 | f"auto_functionalized as auto_functionalized\n" |
| 153 | f"from torch._inductor.pattern_matcher import *\n" |
| 154 | f"vllm = torch.ops.vllm", |
| 155 | "vllm_ir = torch.ops.vllm_ir", |
| 156 | file=f, |
| 157 | ) |
| 158 | |
| 159 | for node, patterns in pm_pass.patterns.items(): |
| 160 | # fix the operator.getitem repr |
| 161 | if node[1] == operator.getitem: |
| 162 | node_repr = f"({repr(node[0])}, operator.getitem)" |
| 163 | else: |
| 164 | node_repr = repr(node) |
| 165 | |
| 166 | node_repr = self._replace_op_overloads(node_repr) |
| 167 | |
| 168 | print(f"\n\n# Patterns for op: {node_repr}", file=f) |
| 169 | for i, pattern in enumerate(patterns): |
| 170 | # reserve auto_functionalized ahead of time |
| 171 | pp = PatternPrettyPrinter() |
| 172 | pp.namespace.create_name("auto_functionalized", None) |
| 173 | |
| 174 | # Assemble pattern |
| 175 | out_node = pp.pretty_print(pattern.pattern) |
| 176 | pattern_repr = "\n".join( |
| 177 | [f"def pattern_{i}():"] |
| 178 | + [ |
| 179 | f"{pp.memoized_objs_names[key]} = " |
| 180 | f"{pp.memoized_objs_pp[key]}" |
| 181 | for key in pp.memoized_objs_names |
| 182 | ] |
| 183 | + [f"return {out_node}"] |
| 184 | ).replace("\n", "\n ") |
| 185 | |
| 186 | pattern_repr = self._replace_op_overloads(pattern_repr) |
| 187 | print(f"{pattern_repr}\n", file=f) |
| 188 | |
| 189 | |
| 190 | P = ParamSpec("P") |
| 191 | R = TypeVar("R") |
| 192 | |
| 193 | |
| 194 | class VllmPatternReplacement(ABC, Generic[P, R]): |
| 195 | """A pattern/replacement pair for FX graph fusion. |
| 196 | |
| 197 | Implement the three abstract members below, then pass |
| 198 | instances to VllmFusionPatternMatcherPass.register(). The pass will |
| 199 | find every occurrence of `pattern` in the graph and substitute it |
| 200 | with `replacement`. |
| 201 | """ |
| 202 | |
| 203 | # TODO(Badr): bound methods work for pattern registration since |
| 204 | # PyTorch 2.10. Once vLLM requires torch>=2.11, replace these properties |
| 205 | # with plain methods and drop the closure indirection. |
| 206 | @property |
| 207 | @abstractmethod |
| 208 | def pattern(self) -> Callable[P, R]: |
| 209 | """Returns a closure defining the FX subgraph to search for.""" |
| 210 | ... |
| 211 | |
| 212 | @property |
| 213 | @abstractmethod |
| 214 | def replacement(self) -> Callable[P, R]: |
| 215 | """Returns a closure defining the FX subgraph to |
| 216 | substitute in place of each match. |
| 217 | """ |
| 218 | ... |
| 219 | |
| 220 | @abstractmethod |
| 221 | def get_inputs(self) -> list[torch.Tensor]: |
| 222 | """Example tensors used to trace pattern and replacement.""" |
| 223 | ... |
| 224 | |
| 225 | # Helpers for get_inputs: uninitialized tensors of common dtypes. |
| 226 | @staticmethod |
| 227 | def empty(*args, **kwargs) -> torch.Tensor: |
| 228 | return torch.empty(*args, device=DEVICE_TYPE, **kwargs) |
| 229 | |
| 230 | @staticmethod |
| 231 | def empty_bf16(*args, **kwargs) -> torch.Tensor: |
| 232 | return torch.empty(*args, dtype=torch.bfloat16, device=DEVICE_TYPE, **kwargs) |
| 233 | |
| 234 | @staticmethod |
| 235 | def empty_fp16(*args, **kwargs) -> torch.Tensor: |
| 236 | return torch.empty(*args, dtype=torch.float16, device=DEVICE_TYPE, **kwargs) |
| 237 | |
| 238 | @staticmethod |
| 239 | def empty_fp32(*args, **kwargs) -> torch.Tensor: |
| 240 | return torch.empty(*args, dtype=torch.float32, device=DEVICE_TYPE, **kwargs) |
| 241 | |
| 242 | @staticmethod |
| 243 | def empty_i32(*args, **kwargs) -> torch.Tensor: |
| 244 | return torch.empty(*args, dtype=torch.int32, device=DEVICE_TYPE, **kwargs) |
| 245 | |
| 246 | |
| 247 | def _fx_view_to_reshape(gm: fx.GraphModule) -> None: |
| 248 | from torch._inductor.fx_passes.post_grad import view_to_reshape |
| 249 | |
| 250 | view_to_reshape(gm) |
| 251 | |
| 252 | |
| 253 | def fold_consecutive_reshapes(gm: fx.GraphModule) -> None: |
| 254 | """Fold consecutive reshape ops into a single reshape. |
| 255 | |
| 256 | ``make_fx`` faithfully records every view/reshape the Python code performs, |
| 257 | so patterns like ``x.reshape(a, b).reshape(c, d)`` produce two reshape |
| 258 | nodes. Inductor's own optimisation would fold these, but |
| 259 | ``pm.register_replacement``'s ``trace_fn`` runs before Inductor, so we |
| 260 | must fold them ourselves for the pattern to match the compiled graph. |
| 261 | |
| 262 | When reshape(A, shape1) feeds only into reshape(result, shape2), |
| 263 | the first reshape is redundant -- replace with reshape(A, shape2). |
| 264 | """ |
| 265 | aten_reshape = torch.ops.aten.reshape.default |
| 266 | for node in list(gm.graph.nodes): |
| 267 | if not is_func(node, aten_reshape): |
| 268 | continue |
| 269 | inp = node.args[0] |
| 270 | if not isinstance(inp, fx.Node) or not is_func(inp, aten_reshape): |
| 271 | continue |
| 272 | if len(inp.users) != 1: |
| 273 | continue |
| 274 | original_input = inp.args[0] |
| 275 | node.args = (original_input, node.args[1]) |
| 276 | inp.replace_all_uses_with(original_input) |
| 277 | gm.graph.erase_node(inp) |
| 278 | |
| 279 | |
| 280 | def remove_noop_reshapes(gm: fx.GraphModule) -> None: |
| 281 | """Drop reshape ops whose output shape equals their input shape. |
| 282 | |
| 283 | Companion to :func:`fold_consecutive_reshapes`. ``make_fx`` records a |
| 284 | reshape to the shape the input already has; the compiled graph has already |
| 285 | dropped it, so the pattern only matches once we drop it too. |
| 286 | """ |
| 287 | aten_reshape = torch.ops.aten.reshape.default |
| 288 | for node in list(gm.graph.nodes): |
| 289 | if not is_func(node, aten_reshape): |
| 290 | continue |
| 291 | inp = node.args[0] |
| 292 | if not isinstance(inp, fx.Node): |
| 293 | continue |
| 294 | val, inp_val = node.meta.get("val"), inp.meta.get("val") |
| 295 | if val is None or inp_val is None or val.shape != inp_val.shape: |
| 296 | continue |
| 297 | node.replace_all_uses_with(inp) |
| 298 | gm.graph.erase_node(node) |
| 299 | |
| 300 | |
| 301 | def _remove_noop_permutes(gm: fx.GraphModule) -> None: |
| 302 | for node in gm.graph.nodes: |
| 303 | if not is_func(node, torch.ops.aten.permute.default): |
| 304 | continue |
| 305 | dims = node.args[1] |
| 306 | if any(dim != i for i, dim in enumerate(dims)): |
| 307 | continue |
| 308 | node.replace_all_uses_with(node.args[0]) |
| 309 | gm.graph.erase_node(node) |
| 310 | |
| 311 | |
| 312 | class VllmFusionPatternMatcherPass(VllmPatternMatcherPass): |
| 313 | """A VllmPatternMatcherPass for passes that use VllmPatternReplacement objects. |
| 314 | Subclasses register patterns via self.register() in their own __init__. |
| 315 | """ |
| 316 | |
| 317 | def __init__(self, config: VllmConfig, pass_name: str) -> None: |
| 318 | super().__init__(config) |
| 319 | self.pass_name = pass_name |
| 320 | self.pm_pass = PatternMatcherPass(pass_name=pass_name) |
| 321 | self._pattern_replacements: list[VllmPatternReplacement] = [] |
| 322 | |
| 323 | @enable_fake_mode |
| 324 | def register(self, pr: VllmPatternReplacement) -> None: |
| 325 | pm.register_replacement( |
| 326 | pr.pattern, |
| 327 | pr.replacement, |
| 328 | pr.get_inputs(), |
| 329 | self._trace_fn, |
| 330 | self.pm_pass, |
| 331 | ) |
| 332 | self._pattern_replacements.append(pr) |
| 333 | |
| 334 | def uuid(self) -> str: |
| 335 | return VllmInductorPass.hash_source( |
| 336 | type(self), |
| 337 | *[type(pr) for pr in self._pattern_replacements], |
| 338 | ) |
| 339 | |
| 340 | @staticmethod |
| 341 | def _trace_fn(*args: Any, **kwargs: Any) -> fx.GraphModule: |
| 342 | gm = pm.fwd_only(*args, **kwargs) |
| 343 | _fx_view_to_reshape(gm) |
| 344 | _remove_noop_permutes(gm) |
| 345 | return gm |
| 346 | |
| 347 | @VllmInductorPass.time_and_log |
| 348 | def __call__(self, graph: torch.fx.Graph) -> None: |
| 349 | self.matched_count = self.pm_pass.apply(graph) |
| 350 | VllmPatternMatcherPass.match_table[self.pass_name] += self.matched_count |