| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import ast |
| 5 | import dataclasses |
| 6 | import hashlib |
| 7 | import json |
| 8 | import operator |
| 9 | import os |
| 10 | import pprint |
| 11 | import time |
| 12 | from collections import defaultdict |
| 13 | from collections.abc import Callable, Generator, Sequence |
| 14 | from contextlib import contextmanager |
| 15 | from copy import deepcopy |
| 16 | from functools import partial |
| 17 | from typing import Any |
| 18 | |
| 19 | import torch |
| 20 | import torch.fx as fx |
| 21 | from torch._dynamo.utils import dynamo_timed |
| 22 | from torch._logging._internal import trace_structured |
| 23 | from torch.fx._lazy_graph_module import _use_lazy_graph_module |
| 24 | |
| 25 | import vllm.envs as envs |
| 26 | from vllm.compilation.codegen import ( |
| 27 | compile_execution_fn, |
| 28 | generate_execution_code, |
| 29 | ) |
| 30 | from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig |
| 31 | from vllm.config.compilation import DynamicShapesType |
| 32 | from vllm.config.utils import Range, hash_factors |
| 33 | from vllm.logger import init_logger |
| 34 | from vllm.logging_utils import lazy |
| 35 | from vllm.platforms import current_platform |
| 36 | from vllm.tracing import instrument, instrument_manual |
| 37 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 38 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 39 | |
| 40 | from .compiler_interface import ( |
| 41 | CompilerInterface, |
| 42 | EagerAdaptor, |
| 43 | InductorAdaptor, |
| 44 | InductorStandaloneAdaptor, |
| 45 | is_compile_cache_enabled, |
| 46 | ) |
| 47 | from .counter import compilation_counter |
| 48 | from .partition_rules import ( |
| 49 | inductor_partition_rule_context, |
| 50 | should_split, |
| 51 | ) |
| 52 | from .passes.inductor_pass import InductorPass, pass_context |
| 53 | from .passes.ir.inplace_functionalization import VllmIRInplaceFunctionalizationPass |
| 54 | from .passes.pass_manager import PostGradPassManager |
| 55 | |
| 56 | logger = init_logger(__name__) |
| 57 | |
| 58 | |
| 59 | def make_copy_and_call( |
| 60 | sym_tensor_indices: list[int], |
| 61 | input_buffers: list[torch.Tensor | None], |
| 62 | callable_fn: Callable[..., Any], |
| 63 | ) -> Callable[..., Any]: |
| 64 | """Create a wrapper that copies inputs to static buffers before calling. |
| 65 | |
| 66 | This is used for cudagraph input copying where we need to copy dynamic |
| 67 | tensors to static buffers before invoking the compiled graph. |
| 68 | |
| 69 | Args: |
| 70 | sym_tensor_indices: Indices of tensors with symbolic shapes |
| 71 | input_buffers: List of static buffers (can contain None for lazy init) |
| 72 | callable_fn: The compiled function to call |
| 73 | |
| 74 | Returns: |
| 75 | A wrapper function that copies inputs and calls the compiled function |
| 76 | |
| 77 | """ |
| 78 | |
| 79 | def copy_and_call(*args: Any) -> Any: |
| 80 | list_args = list(args) |
| 81 | for i, index in enumerate(sym_tensor_indices): |
| 82 | runtime_tensor = list_args[index] |
| 83 | runtime_shape = runtime_tensor.shape[0] |
| 84 | |
| 85 | # lazy initialization of buffer on first call |
| 86 | if input_buffers[i] is None: |
| 87 | input_buffers[i] = runtime_tensor.clone() |
| 88 | |
| 89 | static_tensor = input_buffers[i][:runtime_shape] # type: ignore[index] |
| 90 | static_tensor.copy_(runtime_tensor) |
| 91 | list_args[index] = static_tensor |
| 92 | return callable_fn(*list_args) |
| 93 | |
| 94 | return copy_and_call |
| 95 | |
| 96 | |
| 97 | def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface: |
| 98 | assert not envs.VLLM_USE_MEGA_AOT_ARTIFACT or envs.VLLM_USE_STANDALONE_COMPILE, ( |
| 99 | "VLLM_USE_MEGA_AOT_ARTIFACT=1 requires VLLM_USE_STANDALONE_COMPILE=1" |
| 100 | ) |
| 101 | |
| 102 | if compilation_config.backend == "inductor": |
| 103 | # Use standalone compile only if requested, version is new enough, |
| 104 | # and the symbol actually exists in this PyTorch build. |
| 105 | if envs.VLLM_USE_STANDALONE_COMPILE and hasattr( |
| 106 | torch._inductor, "standalone_compile" |
| 107 | ): |
| 108 | logger.debug("Using InductorStandaloneAdaptor") |
| 109 | return InductorStandaloneAdaptor( |
| 110 | compilation_config.compile_cache_save_format |
| 111 | ) |
| 112 | else: |
| 113 | logger.debug("Using InductorAdaptor") |
| 114 | return InductorAdaptor() |
| 115 | elif compilation_config.backend == "eager": |
| 116 | logger.debug("Using EagerAdaptor") |
| 117 | return EagerAdaptor() |
| 118 | else: |
| 119 | logger.debug("Using custom backend: %s", compilation_config.backend) |
| 120 | compiler = resolve_obj_by_qualname(current_platform.get_compile_backend())() |
| 121 | assert isinstance(compiler, CompilerInterface) |
| 122 | return compiler |
| 123 | |
| 124 | |
| 125 | class CompilerManager: |
| 126 | """A manager to manage the compilation process, including |
| 127 | caching the compiled graph, loading the compiled graph, |
| 128 | and compiling the graph. |
| 129 | |
| 130 | The cache is a dict mapping |
| 131 | `(runtime_shape, graph_index, backend_name)` |
| 132 | to `any_data` returned from the compiler. |
| 133 | |
| 134 | When serializing the cache, we save it to a Python file |
| 135 | for readability. We don't use json here because json doesn't |
| 136 | support int as key. |
| 137 | """ |
| 138 | |
| 139 | def __init__(self, compilation_config: CompilationConfig) -> None: |
| 140 | self.cache: dict[tuple[Range, int, str], Any] = dict() |
| 141 | self.is_cache_updated = False |
| 142 | self.compilation_config = compilation_config |
| 143 | self.compiler = make_compiler(compilation_config) |
| 144 | self.loaded_artifacts: dict[str, Any] = {} |
| 145 | self.prefix: str = "" |
| 146 | |
| 147 | def _log_prefix(self) -> str: |
| 148 | return f"[{self.prefix}] " if self.prefix else "" |
| 149 | |
| 150 | def compute_hash(self, vllm_config: VllmConfig) -> str: |
| 151 | return self.compiler.compute_hash(vllm_config) |
| 152 | |
| 153 | @contextmanager |
| 154 | def compile_context(self, compile_range: Range) -> Generator[None, None, None]: |
| 155 | """Provide compilation context for the duration of compilation to set |
| 156 | any torch global properties we want to scope to a single Inductor |
| 157 | compilation (e.g. partition rules, pass context).""" |
| 158 | with pass_context(compile_range): |
| 159 | if self.compilation_config.use_inductor_graph_partition: |
| 160 | with inductor_partition_rule_context( |
| 161 | self.compilation_config.splitting_ops |
| 162 | ): |
| 163 | yield |
| 164 | else: |
| 165 | yield |
| 166 | |
| 167 | def initialize_cache( |
| 168 | self, cache_dir: str, disable_cache: bool = False, prefix: str = "" |
| 169 | ) -> None: |
| 170 | """Initialize the cache directory for the compiler. |
| 171 | |
| 172 | The organization of the cache directory is as follows: |
| 173 | cache_dir=/path/to/hash_str/rank_i_j/prefix/ |
| 174 | inside cache_dir, there will be: |
| 175 | - vllm_compile_cache.py |
| 176 | - computation_graph.py |
| 177 | - transformed_code.py |
| 178 | |
| 179 | for multiple prefixes, they can share the same |
| 180 | base cache dir of /path/to/hash_str/rank_i_j/ , |
| 181 | to store some common compilation artifacts. |
| 182 | """ |
| 183 | self.disable_cache = disable_cache |
| 184 | self.cache_dir = cache_dir |
| 185 | self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py") |
| 186 | self.prefix = prefix |
| 187 | |
| 188 | if not disable_cache and os.path.exists(self.cache_file_path): |
| 189 | # load the cache from the file |
| 190 | with open(self.cache_file_path) as f: |
| 191 | # we use ast.literal_eval to parse the data |
| 192 | # because it is a safe way to parse Python literals. |
| 193 | # do not use eval(), it is unsafe. |
| 194 | cache = ast.literal_eval(f.read()) |
| 195 | |
| 196 | def check_type(value: Any, ty: type) -> None: |
| 197 | if not isinstance(value, ty): |
| 198 | raise TypeError(f"Expected {ty} but got {type(value)} for {value}") |
| 199 | |
| 200 | def parse_key(key: Any) -> tuple[Range, int, str]: |
| 201 | range_tuple, graph_index, compiler_name = key |
| 202 | check_type(graph_index, int) |
| 203 | check_type(compiler_name, str) |
| 204 | if isinstance(range_tuple, tuple): |
| 205 | start, end = range_tuple |
| 206 | check_type(start, int) |
| 207 | check_type(end, int) |
| 208 | range_tuple = Range(start=start, end=end) |
| 209 | check_type(range_tuple, Range) |
| 210 | return range_tuple, graph_index, compiler_name |
| 211 | |
| 212 | self.cache = {parse_key(key): value for key, value in cache.items()} |
| 213 | |
| 214 | self.compiler.initialize_cache( |
| 215 | cache_dir=cache_dir, disable_cache=disable_cache, prefix=prefix |
| 216 | ) |
| 217 | |
| 218 | def save_to_file(self) -> None: |
| 219 | if self.disable_cache or not self.is_cache_updated: |
| 220 | return |
| 221 | printer = pprint.PrettyPrinter(indent=4) |
| 222 | data = printer.pformat(self.cache) |
| 223 | with open(self.cache_file_path, "w") as f: |
| 224 | f.write(data) |
| 225 | |
| 226 | def load( |
| 227 | self, |
| 228 | graph: fx.GraphModule, |
| 229 | example_inputs: list[Any], |
| 230 | graph_index: int, |
| 231 | compile_range: Range, |
| 232 | ) -> Callable[..., Any] | None: |
| 233 | if (compile_range, graph_index, self.compiler.name) not in self.cache: |
| 234 | return None |
| 235 | |
| 236 | def parse_value(value: Any) -> tuple[tuple[str, str], str]: |
| 237 | assert isinstance(value, dict) |
| 238 | handle = value["graph_handle"] |
| 239 | assert isinstance(handle[0], str) |
| 240 | assert isinstance(handle[1], str) |
| 241 | cache_key = value["cache_key"] |
| 242 | return handle, cache_key |
| 243 | |
| 244 | try: |
| 245 | handle, cache_key = parse_value( |
| 246 | self.cache[(compile_range, graph_index, self.compiler.name)] |
| 247 | ) |
| 248 | except Exception: |
| 249 | # When the cache is outdated, we should ignore the existing file. |
| 250 | # This should cause the correct cache to be generated again. |
| 251 | return None |
| 252 | |
| 253 | compiled_graph = self.compiler.load( |
| 254 | handle, graph, example_inputs, graph_index, compile_range |
| 255 | ) |
| 256 | self.loaded_artifacts[cache_key] = compiled_graph |
| 257 | logger.debug( |
| 258 | "Directly load the %s-th graph for compile range %sfrom %s via handle %s", |
| 259 | graph_index, |
| 260 | str(compile_range), |
| 261 | self.compiler.name, |
| 262 | handle, |
| 263 | ) |
| 264 | return compiled_graph |
| 265 | |
| 266 | @instrument(span_name="Compile graph") |
| 267 | def compile( |
| 268 | self, |
| 269 | graph: fx.GraphModule, |
| 270 | example_inputs: list[Any], |
| 271 | additional_inductor_config: dict[str, Any], |
| 272 | compilation_config: CompilationConfig, |
| 273 | compile_range: Range, |
| 274 | graph_index: int = 0, |
| 275 | num_graphs: int = 1, |
| 276 | is_encoder: bool = False, |
| 277 | ) -> Any: |
| 278 | if graph_index == 0: |
| 279 | # before compiling the first graph, record the start time |
| 280 | global compilation_start_time |
| 281 | compilation_start_time = time.perf_counter() |
| 282 | |
| 283 | compilation_counter.num_backend_compilations += 1 |
| 284 | |
| 285 | compiled_graph = None |
| 286 | handle = None |
| 287 | |
| 288 | # try to load from the cache |
| 289 | compiled_graph = self.load(graph, example_inputs, graph_index, compile_range) |
| 290 | if compiled_graph is not None: |
| 291 | if graph_index == num_graphs - 1: |
| 292 | # after loading the last graph for this shape, record the time. |
| 293 | # there can be multiple graphs due to piecewise compilation. |
| 294 | elapsed = time.perf_counter() - compilation_start_time |
| 295 | logger.info( |
| 296 | "%sDirectly load the compiled graph(s) for compile " |
| 297 | "range %s from the cache, took %.3f s", |
| 298 | self._log_prefix(), |
| 299 | str(compile_range), |
| 300 | elapsed, |
| 301 | ) |
| 302 | return compiled_graph |
| 303 | |
| 304 | # no compiler cached the graph, or the cache is disabled, |
| 305 | # we need to compile it |
| 306 | if isinstance(self.compiler, InductorAdaptor): |
| 307 | # Let compile_fx generate a key for us |
| 308 | maybe_key = None |
| 309 | else: |
| 310 | maybe_key = "artifact_compile_range_" |
| 311 | maybe_key += f"{compile_range.start}_{compile_range.end}" |
| 312 | maybe_key += f"_subgraph_{graph_index}" |
| 313 | with self.compile_context(compile_range): |
| 314 | # There is a compilation time optimization here. |
| 315 | # |
| 316 | # If the (input metadata, graph, compiler config) are the same, then |
| 317 | # we want to avoid compiling the same artifact again. If we didn't |
| 318 | # do this optimization, the backend compilation (InductorAdaptor or |
| 319 | # InductorStandaloneAdaptor) |
| 320 | # is able to cache hit and produce an artifact faster if it was |
| 321 | # already created, but it is still a duplicate artifact that |
| 322 | # requires unnecessary things e.g. disk IO. |
| 323 | # |
| 324 | # The optimization is: If the backend compilation cache hits, |
| 325 | # then do an early return from the backend compilation and look up |
| 326 | # which of the previous in-memory artifacts we created to reuse. |
| 327 | # |
| 328 | # We implemented this by monkey-patching torch (torch does not |
| 329 | # easily expose the cache_key function), but in the future torch |
| 330 | # should expose the cache_key function that we can just call |
| 331 | # directly before invoking backend compilation. |
| 332 | cache_key = None |
| 333 | orig = torch._functorch._aot_autograd.autograd_cache.autograd_cache_key |
| 334 | |
| 335 | def autograd_cache_key(*args, **kwargs): |
| 336 | result = orig(*args, **kwargs) |
| 337 | if result is None: |
| 338 | return None |
| 339 | nonlocal cache_key |
| 340 | cache_key = result[0] |
| 341 | if cache_key in self.loaded_artifacts: |
| 342 | raise StopCompiling() |
| 343 | return result |
| 344 | |
| 345 | from unittest.mock import patch |
| 346 | |
| 347 | with ( |
| 348 | # Graphs that are isometric (different node names but same |
| 349 | # structure) should be treated as the same. |
| 350 | torch._functorch.config.patch(autograd_cache_normalize_inputs=True), |
| 351 | patch( |
| 352 | "torch._functorch._aot_autograd.autograd_cache.autograd_cache_key", |
| 353 | autograd_cache_key, |
| 354 | ), |
| 355 | ): |
| 356 | try: |
| 357 | compiled_graph, handle = self.compiler.compile( |
| 358 | graph, |
| 359 | example_inputs, |
| 360 | additional_inductor_config, |
| 361 | compile_range, |
| 362 | maybe_key, |
| 363 | ) |
| 364 | except StopCompiling: |
| 365 | assert cache_key is not None |
| 366 | compiled_graph = self.loaded_artifacts[cache_key] |
| 367 | if cache_key is not None and compiled_graph is not None: |
| 368 | self.loaded_artifacts[cache_key] = compiled_graph |
| 369 | |
| 370 | assert compiled_graph is not None, "Failed to compile the graph" |
| 371 | |
| 372 | # store the artifact in the cache |
| 373 | if is_compile_cache_enabled(additional_inductor_config) and handle is not None: |
| 374 | self.cache[(compile_range, graph_index, self.compiler.name)] = { |
| 375 | "graph_handle": handle, |
| 376 | "cache_key": cache_key, |
| 377 | } |
| 378 | compilation_counter.num_cache_entries_updated += 1 |
| 379 | self.is_cache_updated = True |
| 380 | if graph_index == 0: |
| 381 | # adds some info logging for the first graph |
| 382 | logger.info( |
| 383 | "%sCache the graph of compile range %s for later use", |
| 384 | self._log_prefix(), |
| 385 | str(compile_range), |
| 386 | ) |
| 387 | logger.debug_once( |
| 388 | "Store the %s-th graph for compile range%s from %s via handle %s", |
| 389 | graph_index, |
| 390 | str(compile_range), |
| 391 | self.compiler.name, |
| 392 | handle, |
| 393 | ) |
| 394 | |
| 395 | # after compiling the last graph, record the end time |
| 396 | if graph_index == num_graphs - 1: |
| 397 | elapsed = time.perf_counter() - compilation_start_time |
| 398 | logger.info( |
| 399 | "%sCompiling a graph for compile range %s takes %.2f s", |
| 400 | self._log_prefix(), |
| 401 | str(compile_range), |
| 402 | elapsed, |
| 403 | ) |
| 404 | |
| 405 | return compiled_graph |
| 406 | |
| 407 | |
| 408 | class StopCompiling(BaseException): |
| 409 | pass |
| 410 | |
| 411 | |
| 412 | @dataclasses.dataclass |
| 413 | class SplitItem: |
| 414 | submod_name: str |
| 415 | graph_id: int |
| 416 | is_splitting_graph: bool |
| 417 | graph: fx.GraphModule |
| 418 | |
| 419 | |
| 420 | def _is_empty_allocation_node(node: fx.Node) -> bool: |
| 421 | if node.op == "call_method": |
| 422 | return node.target == "new_empty" |
| 423 | |
| 424 | if node.op != "call_function": |
| 425 | return False |
| 426 | |
| 427 | target = node.target |
| 428 | if target in (torch.empty, torch.empty_like, torch.empty_strided): |
| 429 | return True |
| 430 | |
| 431 | if isinstance(target, torch._ops.OpOverloadPacket): |
| 432 | packet_name = target._qualified_op_name |
| 433 | elif isinstance(target, torch._ops.OpOverload): |
| 434 | packet_name = target.name() |
| 435 | else: |
| 436 | return False |
| 437 | |
| 438 | return packet_name.startswith("aten::empty") or packet_name.startswith( |
| 439 | "aten::new_empty" |
| 440 | ) |
| 441 | |
| 442 | |
| 443 | def _merge_empty_only_subgraphs( |
| 444 | node_to_subgraph_id: dict[fx.Node, int], |
| 445 | split_op_graphs: list[int], |
| 446 | ) -> None: |
| 447 | """Merge a partition that only contains an empty allocation op into the |
| 448 | previous partition. This avoids generating standalone empty submodules, |
| 449 | which can lead to empty cudagraph captures. |
| 450 | """ |
| 451 | nodes_by_subgraph_id: dict[int, list[fx.Node]] = defaultdict(list) |
| 452 | for node, subgraph_id in node_to_subgraph_id.items(): |
| 453 | nodes_by_subgraph_id[subgraph_id].append(node) |
| 454 | |
| 455 | splitting_subgraphs = set(split_op_graphs) |
| 456 | prev_non_splitting_subgraph_id: int | None = None |
| 457 | |
| 458 | max_subgraph_id = max(node_to_subgraph_id.values(), default=-1) |
| 459 | for subgraph_id in range(max_subgraph_id + 1): |
| 460 | nodes = nodes_by_subgraph_id.get(subgraph_id, []) |
| 461 | if not nodes: |
| 462 | continue |
| 463 | |
| 464 | is_non_splitting_subgraph = subgraph_id not in splitting_subgraphs |
| 465 | is_empty_only_subgraph = len(nodes) == 1 and _is_empty_allocation_node(nodes[0]) |
| 466 | merged = False |
| 467 | |
| 468 | if is_empty_only_subgraph and prev_non_splitting_subgraph_id is not None: |
| 469 | # Safety check: don't move allocation before any input producer. |
| 470 | empty_node = nodes[0] |
| 471 | if all( |
| 472 | input_node.op == "placeholder" |
| 473 | or node_to_subgraph_id[input_node] <= prev_non_splitting_subgraph_id |
| 474 | for input_node in empty_node.all_input_nodes |
| 475 | ): |
| 476 | node_to_subgraph_id[empty_node] = prev_non_splitting_subgraph_id |
| 477 | merged = True |
| 478 | |
| 479 | if not merged and is_non_splitting_subgraph: |
| 480 | prev_non_splitting_subgraph_id = subgraph_id |
| 481 | |
| 482 | |
| 483 | def _decompose_size_nodes(graph: fx.GraphModule) -> None: |
| 484 | """Decompose x.size() into per-dim sym_size.int calls. |
| 485 | |
| 486 | torch.Size objects cannot cross split boundaries because aot_autograd |
| 487 | cannot handle them as submodule outputs. This replaces each size() call |
| 488 | with individual sym_size.int(x, dim) nodes: |
| 489 | - Dynamic dims (SymInt) → new sym_size.int node |
| 490 | - Static dims (plain int) → inlined as literal constant |
| 491 | """ |
| 492 | # Dynamo captures x.size()/x.shape as call_method target="size". |
| 493 | size_nodes = list(graph.graph.find_nodes(op="call_method", target="size")) |
| 494 | |
| 495 | for node in size_nodes: |
| 496 | # Only x.size() (no dim) returns a torch.Size tuple that can't cross |
| 497 | # split boundaries. x.size(dim) already returns a scalar SymInt/int, |
| 498 | # which crosses fine, so leave it untouched. |
| 499 | if len(node.args) > 1 or "dim" in node.kwargs: |
| 500 | continue |
| 501 | tensor_node = node.args[0] |
| 502 | ev = tensor_node.meta.get("example_value") |
| 503 | assert ev is not None, ( |
| 504 | f"Tensor node '{tensor_node.name}' has no example_value metadata. " |
| 505 | f"Cannot decompose size node '{node.name}'." |
| 506 | ) |
| 507 | |
| 508 | # Build per-dim replacements: sym_size.int node or literal int. |
| 509 | dims: list[fx.Node | int] = [] |
| 510 | with graph.graph.inserting_after(tensor_node): |
| 511 | for i in range(ev.dim()): |
| 512 | dim_val = ev.shape[i] |
| 513 | if isinstance(dim_val, torch.SymInt): |
| 514 | dn = graph.graph.call_function( |
| 515 | torch.ops.aten.sym_size.int, args=(tensor_node, i) |
| 516 | ) |
| 517 | dn.meta["example_value"] = dim_val |
| 518 | dims.append(dn) |
| 519 | elif isinstance(dim_val, int): |
| 520 | dims.append(dim_val) |
| 521 | else: |
| 522 | raise AssertionError( |
| 523 | f"dim_val is either torch.SymInt or int, " |
| 524 | f"got {type(dim_val)} for dim {i} of " |
| 525 | f"'{node.name}'" |
| 526 | ) |
| 527 | |
| 528 | # Replace size node in each user's args. |
| 529 | for user in list(node.users): |
| 530 | if ( |
| 531 | user.op == "call_function" |
| 532 | and user.target is operator.getitem |
| 533 | and len(user.args) == 2 |
| 534 | and user.args[0] is node |
| 535 | ): |
| 536 | # getitem(size, idx) → replace with dims[idx] directly. |
| 537 | idx = user.args[1] |
| 538 | assert isinstance(idx, int), ( |
| 539 | f"Expected literal int index for getitem on size(), " |
| 540 | f"got {type(idx).__name__}: {idx}" |
| 541 | ) |
| 542 | user.replace_all_uses_with(dims[idx]) |
| 543 | graph.graph.erase_node(user) |
| 544 | else: |
| 545 | # User consumes the full size tuple (e.g. view(clone, size)) |
| 546 | # → view(clone, d0, d1, ...) |
| 547 | new_args = [] |
| 548 | for arg in user.args: |
| 549 | if arg is node: |
| 550 | new_args.extend(dims) |
| 551 | else: |
| 552 | new_args.append(arg) |
| 553 | user.args = tuple(new_args) |
| 554 | graph.graph.erase_node(node) |
| 555 | |
| 556 | |
| 557 | def split_graph( |
| 558 | graph: fx.GraphModule, splitting_ops: list[str] |
| 559 | ) -> tuple[fx.GraphModule, list[SplitItem]]: |
| 560 | _decompose_size_nodes(graph) |
| 561 | |
| 562 | # split graph by ops |
| 563 | subgraph_id = 0 |
| 564 | node_to_subgraph_id: dict[fx.Node, int] = {} |
| 565 | split_op_graphs: list[int] = [] |
| 566 | for node in graph.graph.nodes: |
| 567 | if node.op in ("output", "placeholder"): |
| 568 | continue |
| 569 | |
| 570 | # Check if this is a getitem operation on a node from an earlier subgraph. |
| 571 | # If so, assign it to the same subgraph as its input to avoid passing entire |
| 572 | # tuple as input to submodules, which is against standalone_compile and |
| 573 | # AoTAutograd input requirement. |
| 574 | if node.op == "call_function" and node.target == operator.getitem: |
| 575 | # Assign this getitem to the same subgraph as its input |
| 576 | input_node = node.args[0] |
| 577 | if input_node.op != "placeholder": |
| 578 | assert input_node in node_to_subgraph_id |
| 579 | node_to_subgraph_id[node] = node_to_subgraph_id[input_node] |
| 580 | continue |
| 581 | |
| 582 | if should_split(node, splitting_ops): |
| 583 | subgraph_id += 1 |
| 584 | node_to_subgraph_id[node] = subgraph_id |
| 585 | split_op_graphs.append(subgraph_id) |
| 586 | |
| 587 | # keep consecutive splitting ops together |
| 588 | # (we know node.next exists because node isn't the last (output) node) |
| 589 | if should_split(node.next, splitting_ops): |
| 590 | # this will get incremented by the next node |
| 591 | subgraph_id -= 1 |
| 592 | else: |
| 593 | subgraph_id += 1 |
| 594 | else: |
| 595 | node_to_subgraph_id[node] = subgraph_id |
| 596 | |
| 597 | _merge_empty_only_subgraphs(node_to_subgraph_id, split_op_graphs) |
| 598 | |
| 599 | # `keep_original_order` is important! |
| 600 | # otherwise pytorch might reorder the nodes and |
| 601 | # the semantics of the graph will change when we |
| 602 | # have mutations in the graph |
| 603 | with _use_lazy_graph_module(True): |
| 604 | has_tuple_return = is_torch_equal_or_newer("2.12.0.dev") |
| 605 | tuple_return_kwarg = {"tuple_return": True} if has_tuple_return else {} |
| 606 | split_gm = torch.fx.passes.split_module.split_module( |
| 607 | graph, |
| 608 | None, |
| 609 | lambda node: node_to_subgraph_id[node], |
| 610 | keep_original_order=True, |
| 611 | **tuple_return_kwarg, |
| 612 | ) |
| 613 | |
| 614 | outputs = [] |
| 615 | |
| 616 | names = [name for (name, module) in split_gm.named_modules()] |
| 617 | |
| 618 | for name in names: |
| 619 | if "." in name or name == "": |
| 620 | # recursive child module or the root module |
| 621 | continue |
| 622 | |
| 623 | module = getattr(split_gm, name) |
| 624 | |
| 625 | graph_id = int(name.replace("submod_", "")) |
| 626 | outputs.append(SplitItem(name, graph_id, (graph_id in split_op_graphs), module)) |
| 627 | |
| 628 | # sort by integer graph_id, rather than string name |
| 629 | outputs.sort(key=lambda x: x.graph_id) |
| 630 | |
| 631 | return split_gm, outputs |
| 632 | |
| 633 | |
| 634 | compilation_start_time = 0.0 |
| 635 | |
| 636 | |
| 637 | def wrap_with_cudagraph_if_needed( |
| 638 | piecewise_backend: Any, |
| 639 | vllm_config: VllmConfig, |
| 640 | compilation_config: CompilationConfig, |
| 641 | is_first_graph: bool, |
| 642 | is_last_graph: bool, |
| 643 | ) -> Any: |
| 644 | """Wrap a piecewise backend with CUDA graph wrapper if needed. |
| 645 | This function is shared between VllmBackend and |
| 646 | construct_serializable_fn_from_inductor_cache. |
| 647 | |
| 648 | Args: |
| 649 | piecewise_backend: The backend to wrap |
| 650 | vllm_config: The vLLM configuration |
| 651 | compilation_config: The compilation configuration |
| 652 | is_first_graph: Whether this is the first graph in the sequence |
| 653 | is_last_graph: Whether this is the last graph in the sequence |
| 654 | |
| 655 | Returns: |
| 656 | The wrapped backend if CUDA graphs are enabled, otherwise the original backend |
| 657 | |
| 658 | """ |
| 659 | if ( |
| 660 | not compilation_config.cudagraph_mode.has_piecewise_cudagraphs() |
| 661 | or compilation_config.use_inductor_graph_partition |
| 662 | ): |
| 663 | return piecewise_backend |
| 664 | |
| 665 | # We're using Dynamo-based piecewise splitting, so we wrap |
| 666 | # the whole subgraph with a static graph wrapper. |
| 667 | from .cuda_graph import CUDAGraphOptions |
| 668 | |
| 669 | # resolve the static graph wrapper class (e.g. CUDAGraphWrapper |
| 670 | # class) as platform dependent. |
| 671 | static_graph_wrapper_class = resolve_obj_by_qualname( |
| 672 | current_platform.get_static_graph_wrapper_cls() |
| 673 | ) |
| 674 | |
| 675 | # Always assign PIECEWISE runtime mode to the |
| 676 | # CUDAGraphWrapper for piecewise_backend, to distinguish |
| 677 | # it from the FULL cudagraph runtime mode, no matter it |
| 678 | # is wrapped on a full or piecewise fx graph. |
| 679 | return static_graph_wrapper_class( |
| 680 | runnable=piecewise_backend, |
| 681 | vllm_config=vllm_config, |
| 682 | runtime_mode=CUDAGraphMode.PIECEWISE, |
| 683 | cudagraph_options=CUDAGraphOptions( |
| 684 | debug_log_enable=is_first_graph, |
| 685 | gc_disable=not is_first_graph, |
| 686 | weak_ref_output=is_last_graph, |
| 687 | ), |
| 688 | ) |
| 689 | |
| 690 | |
| 691 | class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc] |
| 692 | """Code adapted from `torch.fx.passes.shape_prop.ShapeProp`. |
| 693 | It runs the given split graph interpreter, and for each submodule in |
| 694 | `compile_submod_names`, creates a PiecewiseBackend and compiles all |
| 695 | ranges up front. |
| 696 | |
| 697 | NOTE: the order in `compile_submod_names` matters, because |
| 698 | it will be used to determine the order of the compiled piecewise |
| 699 | graphs. The first graph will handle logging, and the last graph |
| 700 | has some special cudagraph output handling. |
| 701 | |
| 702 | Note: This class shares similar logic with |
| 703 | reconstruct_serializable_fn_from_mega_artifact in caching.py. |
| 704 | Both create PiecewiseBackend instances and wrap them with cudagraph. |
| 705 | The key difference is: |
| 706 | - reconstruct_serializable_fn_from_mega_artifact: PiecewiseBackend receives |
| 707 | pre-compiled runnables (compiled_runnables is set, graph is None) |
| 708 | - this class: PiecewiseBackend receives the FX graph to compile |
| 709 | (graph is set, compiled_runnables is None) |
| 710 | |
| 711 | |
| 712 | If modifying the backend creation/wrapping logic, consider updating both. |
| 713 | """ |
| 714 | |
| 715 | def __init__( |
| 716 | self, |
| 717 | module: torch.fx.GraphModule, |
| 718 | compile_submod_names: list[str], |
| 719 | vllm_config: VllmConfig, |
| 720 | vllm_backend: "VllmBackend", |
| 721 | ) -> None: |
| 722 | super().__init__(module) |
| 723 | self.compile_submod_names = compile_submod_names |
| 724 | self.compilation_config = vllm_config.compilation_config |
| 725 | self.vllm_config = vllm_config |
| 726 | self.vllm_backend = vllm_backend |
| 727 | # When True, it annoyingly dumps the torch.fx.Graph on errors. |
| 728 | self.extra_traceback = False |
| 729 | |
| 730 | @instrument(span_name="Inductor compilation") |
| 731 | def run(self, *args: Any) -> Any: |
| 732 | return super().run(*args) |
| 733 | |
| 734 | def call_module( |
| 735 | self, |
| 736 | target: torch.fx.node.Target, |
| 737 | args: tuple[torch.fx.node.Argument, ...], |
| 738 | kwargs: dict[str, Any], |
| 739 | ) -> Any: |
| 740 | assert isinstance(target, str) |
| 741 | |
| 742 | gm = getattr(self.module, target) |
| 743 | outputs = gm.graph.output_node().args[0] |
| 744 | output = fx.map_arg(outputs, lambda node: node.meta["example_value"]) |
| 745 | |
| 746 | if target in self.compile_submod_names: |
| 747 | index = self.compile_submod_names.index(target) |
| 748 | submod = self.fetch_attr(target) |
| 749 | |
| 750 | sym_shape_indices = [ |
| 751 | i for i, x in enumerate(args) if isinstance(x, torch.SymInt) |
| 752 | ] |
| 753 | |
| 754 | # Lazy import here to avoid circular import |
| 755 | from torch._inductor.compile_fx import graph_returns_tuple |
| 756 | |
| 757 | from .piecewise_backend import PiecewiseBackend |
| 758 | |
| 759 | piecewise_backend = PiecewiseBackend( |
| 760 | submod, |
| 761 | self.vllm_config, |
| 762 | index, |
| 763 | len(self.compile_submod_names), |
| 764 | sym_shape_indices, |
| 765 | self.vllm_backend, |
| 766 | graph_returns_tuple(submod), |
| 767 | submod_name=target, |
| 768 | ) |
| 769 | |
| 770 | self.module.__dict__[target] = wrap_with_cudagraph_if_needed( |
| 771 | piecewise_backend, |
| 772 | self.vllm_config, |
| 773 | self.compilation_config, |
| 774 | piecewise_backend.is_first_graph, |
| 775 | piecewise_backend.is_last_graph, |
| 776 | ) |
| 777 | |
| 778 | compilation_counter.num_piecewise_capturable_graphs_seen += 1 |
| 779 | |
| 780 | return output |
| 781 | |
| 782 | |
| 783 | # the tag for the part of model being compiled, |
| 784 | # e.g. backbone/eagle_head |
| 785 | model_tag: str = "backbone" |
| 786 | model_is_encoder: bool = False |
| 787 | |
| 788 | |
| 789 | @contextmanager |
| 790 | def set_model_tag(tag: str, is_encoder: bool = False) -> Generator[None, None, None]: |
| 791 | """Context manager to set the model tag.""" |
| 792 | global model_tag |
| 793 | global model_is_encoder |
| 794 | assert tag != model_tag, ( |
| 795 | f"Model tag {tag} is the same as the current tag {model_tag}." |
| 796 | ) |
| 797 | old_tag = model_tag |
| 798 | old_is_encoder = model_is_encoder |
| 799 | |
| 800 | model_tag = tag |
| 801 | model_is_encoder = is_encoder |
| 802 | try: |
| 803 | yield |
| 804 | finally: |
| 805 | model_tag = old_tag |
| 806 | model_is_encoder = old_is_encoder |
| 807 | |
| 808 | |
| 809 | class VllmBackend: |
| 810 | """The compilation backend for `torch.compile` with vLLM. |
| 811 | It is used for compilation mode of `CompilationMode.VLLM_COMPILE`, |
| 812 | where we customize the compilation. |
| 813 | |
| 814 | The major work of this backend is to split the graph into |
| 815 | piecewise graphs, and pass them to the piecewise backend. |
| 816 | |
| 817 | This backend also adds the PostGradPassManager to Inductor config, |
| 818 | which handles the post-grad passes. |
| 819 | """ |
| 820 | |
| 821 | vllm_config: VllmConfig |
| 822 | compilation_config: CompilationConfig |
| 823 | _called: bool = False |
| 824 | # the graph we compiled |
| 825 | graph: fx.GraphModule |
| 826 | # the stiching graph module for all the piecewise graphs |
| 827 | split_gm: fx.GraphModule |
| 828 | piecewise_graphs: list[SplitItem] |
| 829 | returned_callable: Callable[..., Any] |
| 830 | # Inductor passes to run on the graph pre-defunctionalization |
| 831 | post_grad_passes: Sequence[Callable[..., Any]] |
| 832 | compiler_manager: CompilerManager |
| 833 | # Copy of CompilationConfig.inductor_compile_config + |
| 834 | # an entry for PostGradPassManager |
| 835 | inductor_config: dict[str, Any] |
| 836 | |
| 837 | def __init__( |
| 838 | self, |
| 839 | vllm_config: VllmConfig, |
| 840 | prefix: str = "", |
| 841 | is_encoder: bool = False, |
| 842 | ) -> None: |
| 843 | # if the model is initialized with a non-empty prefix, |
| 844 | # then usually it's enough to use that prefix, |
| 845 | # e.g. language_model, vision_model, etc. |
| 846 | # when multiple parts are initialized as independent |
| 847 | # models, we need to use the model_tag to distinguish |
| 848 | # them, e.g. backbone (default), eagle_head, etc. |
| 849 | self.prefix = prefix or model_tag |
| 850 | |
| 851 | # Mark compilation for encoder. |
| 852 | self.is_encoder = is_encoder or model_is_encoder |
| 853 | |
| 854 | # Passes to run on the graph post-grad. |
| 855 | self.pass_manager = resolve_obj_by_qualname( |
| 856 | current_platform.get_pass_manager_cls() |
| 857 | )() |
| 858 | self.pass_key = current_platform.pass_key |
| 859 | |
| 860 | self.vllm_config = vllm_config |
| 861 | self.compilation_config = vllm_config.compilation_config |
| 862 | |
| 863 | self.compiler_manager: CompilerManager = CompilerManager( |
| 864 | self.compilation_config |
| 865 | ) |
| 866 | |
| 867 | # Deepcopy the inductor config to detach the post-grad custom pass |
| 868 | # from CompilationConfig. |
| 869 | # We want to avoid PostGradPassManager in CompilationConfig because |
| 870 | # in future we need PostGradPassManager.uuid() to be executed |
| 871 | # only at compile time. |
| 872 | self.inductor_config = deepcopy(self.compilation_config.inductor_compile_config) |
| 873 | # `torch.compile` is JIT compiled, so we don't need to |
| 874 | # do anything here |
| 875 | |
| 876 | def _log_prefix(self) -> str: |
| 877 | return f"[{self.prefix}] " if self.prefix else "" |
| 878 | |
| 879 | def collect_standalone_compile_artifacts( |
| 880 | self, |
| 881 | ) -> tuple[Any, dict[str, list[int]] | None, dict[str, bool] | None]: |
| 882 | """Collect inductor cache artifacts from all piecewise backends. |
| 883 | |
| 884 | Returns: |
| 885 | tuple: (standalone_compile_artifacts, sym_shape_indices_map, |
| 886 | returns_tuple_map) |
| 887 | - standalone_compile_artifacts: StandaloneCompiledArtifacts |
| 888 | with compiled artifacts |
| 889 | - sym_shape_indices_map: dict mapping submod_name to |
| 890 | sym_shape_indices |
| 891 | - returns_tuple_map: dict mapping submod_name to |
| 892 | returns_tuple |
| 893 | |
| 894 | """ |
| 895 | if not envs.VLLM_USE_MEGA_AOT_ARTIFACT: |
| 896 | return None, None, None |
| 897 | |
| 898 | from .caching import StandaloneCompiledArtifacts |
| 899 | from .piecewise_backend import PiecewiseBackend |
| 900 | |
| 901 | standalone_compile_artifacts = StandaloneCompiledArtifacts() |
| 902 | sym_shape_indices_map = {} |
| 903 | returns_tuple_map = {} |
| 904 | |
| 905 | for name, _ in self.split_gm.named_children(): |
| 906 | # get the actual attribute (shadowed by PiecewiseBackend in __dict__) |
| 907 | child = getattr(self.split_gm, name) |
| 908 | # unwrap the static graph wrapper class if applicable |
| 909 | piecewise_backend = child.runnable if hasattr(child, "runnable") else child |
| 910 | |
| 911 | if not isinstance(piecewise_backend, PiecewiseBackend): |
| 912 | continue |
| 913 | |
| 914 | submod_name = name |
| 915 | sym_shape_indices_map[submod_name] = piecewise_backend.sym_shape_indices |
| 916 | returns_tuple_map[submod_name] = piecewise_backend.returns_tuple |
| 917 | |
| 918 | for shape_str, bytes_data in piecewise_backend.to_bytes().items(): |
| 919 | standalone_compile_artifacts.insert(submod_name, shape_str, bytes_data) |
| 920 | logger.debug( |
| 921 | "collected artifact for %s shape %s (%d bytes)", |
| 922 | submod_name, |
| 923 | shape_str, |
| 924 | len(bytes_data), |
| 925 | ) |
| 926 | |
| 927 | logger.info( |
| 928 | "collected artifacts: %d entries, %d artifacts, %d bytes total", |
| 929 | standalone_compile_artifacts.num_entries(), |
| 930 | standalone_compile_artifacts.num_artifacts(), |
| 931 | standalone_compile_artifacts.size_bytes(), |
| 932 | ) |
| 933 | |
| 934 | logger.debug( |
| 935 | "standalone compile artifact keys: %s", |
| 936 | list(standalone_compile_artifacts.submodule_bytes.keys()), |
| 937 | ) |
| 938 | |
| 939 | return standalone_compile_artifacts, sym_shape_indices_map, returns_tuple_map |
| 940 | |
| 941 | def configure_post_pass(self) -> None: |
| 942 | # TODO proper PassManager? |
| 943 | pre_grad_pass_key = "pre_grad_custom_pass" |
| 944 | assert self.pass_key != pre_grad_pass_key |
| 945 | assert pre_grad_pass_key not in self.inductor_config |
| 946 | self.inductor_config[pre_grad_pass_key] = VllmIRInplaceFunctionalizationPass( |
| 947 | self.vllm_config |
| 948 | ) |
| 949 | |
| 950 | # Make sure pre_grad_custom_pass is not pickled |
| 951 | # as part of AOTAutograd built-in cache key |
| 952 | # TODO(luka) is there a cleaner way to do this |
| 953 | import torch._inductor.config as inductor_config |
| 954 | |
| 955 | ignore = inductor_config._cache_config_ignore_prefix + [pre_grad_pass_key] |
| 956 | assert "_cache_config_ignore_prefix" not in self.inductor_config |
| 957 | self.inductor_config["_cache_config_ignore_prefix"] = ignore |
| 958 | |
| 959 | # Configure the (nominally post-grad) pass manager |
| 960 | self.pass_manager.configure(self.vllm_config) |
| 961 | |
| 962 | # Post-grad custom passes are run using the post_grad_custom_post_pass |
| 963 | # hook. If a pass for that hook exists, add it to the pass manager. |
| 964 | if self.pass_key in self.inductor_config: |
| 965 | if isinstance(self.inductor_config[self.pass_key], PostGradPassManager): |
| 966 | raise ValueError( |
| 967 | "PostGradPassManager can not be kept in CompilationConfig." |
| 968 | ) |
| 969 | else: |
| 970 | # Config should automatically wrap all inductor passes |
| 971 | assert isinstance( |
| 972 | self.compilation_config.inductor_compile_config[self.pass_key], |
| 973 | InductorPass, |
| 974 | ) |
| 975 | self.pass_manager.add( |
| 976 | self.compilation_config.inductor_compile_config[self.pass_key] |
| 977 | ) |
| 978 | self.inductor_config[self.pass_key] = self.pass_manager |
| 979 | |
| 980 | def _log_compilation_config(self): |
| 981 | """Log vLLM compilation config for TORCH_TRACE/tlparse.""" |
| 982 | cc = self.compilation_config |
| 983 | pass_cfg = cc.pass_config |
| 984 | |
| 985 | # Helper to convert lists to comma-separated strings for tlparse display |
| 986 | def list_to_str(lst: list | None) -> str: |
| 987 | if lst is None: |
| 988 | return "" |
| 989 | return ", ".join(str(x) for x in lst) |
| 990 | |
| 991 | # Get enabled passes by introspecting dataclass fields |
| 992 | enabled_passes = [ |
| 993 | f.name |
| 994 | for f in dataclasses.fields(pass_cfg) |
| 995 | if isinstance(getattr(pass_cfg, f.name), bool) and getattr(pass_cfg, f.name) |
| 996 | ] |
| 997 | |
| 998 | trace_structured( |
| 999 | "artifact", |
| 1000 | metadata_fn=lambda: { |
| 1001 | "name": "vllm_compilation_config", |
| 1002 | "encoding": "json", |
| 1003 | }, |
| 1004 | payload_fn=lambda: json.dumps( |
| 1005 | { |
| 1006 | "model": getattr(self.vllm_config.model_config, "model", "unknown"), |
| 1007 | "prefix": self.prefix, |
| 1008 | "mode": str(cc.mode), |
| 1009 | "backend": cc.backend, |
| 1010 | "custom_ops": list_to_str(cc.custom_ops), |
| 1011 | "splitting_ops": list_to_str(cc.splitting_ops), |
| 1012 | "cudagraph_mode": str(cc.cudagraph_mode), |
| 1013 | "compile_sizes": list_to_str(cc.compile_sizes), |
| 1014 | "compile_ranges_endpoints": list_to_str( |
| 1015 | cc.compile_ranges_endpoints |
| 1016 | ), |
| 1017 | "use_inductor_graph_partition": cc.use_inductor_graph_partition, |
| 1018 | "inductor_passes": list_to_str(list(cc.inductor_passes.keys())), |
| 1019 | "enabled_passes": list_to_str(enabled_passes), |
| 1020 | "dynamic_shapes_type": str(cc.dynamic_shapes_config.type), |
| 1021 | "dynamic_shapes_evaluate_guards": cc.dynamic_shapes_config.evaluate_guards, # noqa: E501 |
| 1022 | } |
| 1023 | ), |
| 1024 | ) |
| 1025 | |
| 1026 | @dynamo_timed("vllm_backend") |
| 1027 | def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any: |
| 1028 | from .caching import ( |
| 1029 | VllmSerializableFunction, |
| 1030 | ) |
| 1031 | |
| 1032 | vllm_config = self.vllm_config |
| 1033 | |
| 1034 | self._log_compilation_config() |
| 1035 | |
| 1036 | # Minimal hashing here with existing utilities, reused below. |
| 1037 | |
| 1038 | env_factors = envs.compile_factors() |
| 1039 | env_hash = hash_factors(env_factors) |
| 1040 | # Compute config/compiler/code hashes once and reuse |
| 1041 | config_hash = vllm_config.compute_hash() |
| 1042 | compiler_hash = self.compiler_manager.compute_hash(vllm_config) |
| 1043 | forward_code_files = list(sorted(self.compilation_config.traced_files)) |
| 1044 | |
| 1045 | logger.debug( |
| 1046 | "Traced files (to be considered for compilation cache):\n%s", |
| 1047 | lazy(lambda: "\n".join(forward_code_files)), |
| 1048 | ) |
| 1049 | hash_content = [] |
| 1050 | for filepath in forward_code_files: |
| 1051 | if filepath == "<string>": |
| 1052 | # This means the function was dynamically generated, with |
| 1053 | # e.g. exec(). We can't actually check these. |
| 1054 | continue |
| 1055 | hash_content.append(filepath) |
| 1056 | try: |
| 1057 | with open(filepath) as f: |
| 1058 | hash_content.append(f.read()) |
| 1059 | except (OSError, UnicodeDecodeError): |
| 1060 | logger.warning("Failed to read file %s", filepath) |
| 1061 | continue |
| 1062 | code_hash = hashlib.sha256("\n".join(hash_content).encode()).hexdigest() |
| 1063 | # Clear after consumption |
| 1064 | self.compilation_config.traced_files.clear() |
| 1065 | if not self.compilation_config.cache_dir: |
| 1066 | # no provided cache dir, generate one based on the known factors |
| 1067 | # that affects the compilation. if none of the factors change, |
| 1068 | # the cache dir will be the same so that we can reuse the compiled |
| 1069 | # graph. |
| 1070 | factors = [env_hash, config_hash, code_hash, compiler_hash] |
| 1071 | # Use SHA-256 for cache key hashing to be consistent across |
| 1072 | # compute_hash functions. Truncate for a short cache dir name. |
| 1073 | hash_key = hashlib.sha256(str(factors).encode()).hexdigest()[:10] |
| 1074 | cache_dir = os.path.join( |
| 1075 | envs.VLLM_CACHE_ROOT, "torch_compile_cache", hash_key |
| 1076 | ) |
| 1077 | self.compilation_config.cache_dir = cache_dir |
| 1078 | |
| 1079 | cache_dir = self.compilation_config.cache_dir |
| 1080 | os.makedirs(cache_dir, exist_ok=True) |
| 1081 | self.compilation_config.cache_dir = cache_dir |
| 1082 | rank = vllm_config.parallel_config.rank |
| 1083 | dp_rank = vllm_config.parallel_config.data_parallel_index |
| 1084 | local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}", self.prefix) |
| 1085 | os.makedirs(local_cache_dir, exist_ok=True) |
| 1086 | self.compilation_config.local_cache_dir = local_cache_dir |
| 1087 | |
| 1088 | # Honors opt-outs such as CompilationMode.NONE or VLLM_DISABLE_COMPILE_CACHE. |
| 1089 | disable_cache = not is_compile_cache_enabled(self.inductor_config) |
| 1090 | |
| 1091 | # TODO(patchy): ngram gpu kernel will cause vllm torch compile cache errors. |
| 1092 | is_ngram_gpu_enabled = ( |
| 1093 | vllm_config.speculative_config is not None |
| 1094 | and vllm_config.speculative_config.use_ngram_gpu() |
| 1095 | ) |
| 1096 | disable_cache = disable_cache or is_ngram_gpu_enabled |
| 1097 | |
| 1098 | if disable_cache: |
| 1099 | logger.info("%svLLM's torch.compile cache is disabled.", self._log_prefix()) |
| 1100 | else: |
| 1101 | logger.info( |
| 1102 | "%sUsing cache directory: %s for vLLM's torch.compile", |
| 1103 | self._log_prefix(), |
| 1104 | local_cache_dir, |
| 1105 | ) |
| 1106 | |
| 1107 | self.compiler_manager.initialize_cache( |
| 1108 | local_cache_dir, disable_cache, self.prefix |
| 1109 | ) |
| 1110 | |
| 1111 | # Reuses existing cache key |
| 1112 | |
| 1113 | logger.debug( |
| 1114 | "torch.compile cache factors: env=%s cfg=%s comp=%s code=%s dir=%s", |
| 1115 | env_hash, |
| 1116 | config_hash, |
| 1117 | compiler_hash, |
| 1118 | code_hash, |
| 1119 | local_cache_dir, |
| 1120 | ) |
| 1121 | |
| 1122 | # Persist and log only hash-relevant factors together. |
| 1123 | try: |
| 1124 | logger.debug( |
| 1125 | "Compile env factors (raw):\n%s\nVllm config hash: %s", |
| 1126 | lazy(partial(pprint.pformat, env_factors, width=120)), |
| 1127 | config_hash, |
| 1128 | ) |
| 1129 | meta_path = os.path.join(local_cache_dir, "cache_key_factors.json") |
| 1130 | if not os.path.exists(meta_path): |
| 1131 | with open(meta_path, "w") as f: |
| 1132 | json.dump( |
| 1133 | { |
| 1134 | "env": env_factors, # raw factors used for env_hash |
| 1135 | "config_hash": config_hash, |
| 1136 | "code_hash": code_hash, |
| 1137 | "compiler_hash": compiler_hash, |
| 1138 | }, |
| 1139 | f, |
| 1140 | indent=2, |
| 1141 | sort_keys=True, |
| 1142 | ) |
| 1143 | except Exception: |
| 1144 | # Best-effort only; metadata write failures are non-fatal. |
| 1145 | logger.warning( |
| 1146 | ( |
| 1147 | "Could not write compile cache metadata at %s; continuing without " |
| 1148 | "metadata. Compiled cache remains valid; diagnostics may be " |
| 1149 | "limited." |
| 1150 | ), |
| 1151 | local_cache_dir, |
| 1152 | exc_info=True, |
| 1153 | ) |
| 1154 | |
| 1155 | # when dynamo calls the backend, it means the bytecode |
| 1156 | # transform and analysis are done |
| 1157 | compilation_counter.num_graphs_seen += 1 |
| 1158 | from .monitor import torch_compile_start_time |
| 1159 | |
| 1160 | current_perf = time.perf_counter() |
| 1161 | current_epoch = time.time() |
| 1162 | dynamo_time = current_perf - torch_compile_start_time |
| 1163 | logger.info( |
| 1164 | "%sDynamo bytecode transform time: %.2f s", |
| 1165 | self._log_prefix(), |
| 1166 | dynamo_time, |
| 1167 | ) |
| 1168 | |
| 1169 | # Record Dynamo time in tracing if available |
| 1170 | real_start_time = current_epoch - dynamo_time |
| 1171 | start_time_ns = int(real_start_time * 1e9) |
| 1172 | attributes = {"dynamo.time_seconds": dynamo_time} |
| 1173 | instrument_manual("Dynamo bytecode transform", start_time_ns, None, attributes) |
| 1174 | |
| 1175 | # we control the compilation process, each instance can only be |
| 1176 | # called once |
| 1177 | assert not self._called, "VllmBackend can only be called once" |
| 1178 | |
| 1179 | self.graph = graph |
| 1180 | self.configure_post_pass() |
| 1181 | |
| 1182 | if self.compilation_config.use_inductor_graph_partition: |
| 1183 | # Let Inductor decide partitioning; avoid FX-level pre-splitting. |
| 1184 | fx_split_ops: list[str] = [] |
| 1185 | else: |
| 1186 | fx_split_ops = self.compilation_config.splitting_ops or [] |
| 1187 | |
| 1188 | self.split_gm, self.piecewise_graphs = split_graph(graph, fx_split_ops) |
| 1189 | |
| 1190 | # keep a split_gm copy from BEFORE the interpreter replaces |
| 1191 | # submodules with PiecewiseBackend -- used for serialization |
| 1192 | original_split_gm = None |
| 1193 | if envs.VLLM_USE_MEGA_AOT_ARTIFACT: |
| 1194 | original_split_gm = deepcopy(self.split_gm) |
| 1195 | |
| 1196 | from torch._dynamo.utils import lazy_format_graph_code |
| 1197 | |
| 1198 | # depyf will hook lazy_format_graph_code and dump the graph |
| 1199 | # for debugging, no need to print the graph here |
| 1200 | lazy_format_graph_code("before split", self.graph) |
| 1201 | lazy_format_graph_code("after split", self.split_gm) |
| 1202 | |
| 1203 | # Log the piecewise split graph for TORCH_TRACE/tlparse |
| 1204 | trace_structured( |
| 1205 | "graph_dump", |
| 1206 | metadata_fn=lambda: {"name": "vllm_piecewise_split_graph"}, |
| 1207 | payload_fn=lambda: self.split_gm.print_readable(print_output=False), |
| 1208 | ) |
| 1209 | |
| 1210 | compilation_counter.num_piecewise_graphs_seen += len(self.piecewise_graphs) |
| 1211 | submod_names_to_compile = [ |
| 1212 | item.submod_name |
| 1213 | for item in self.piecewise_graphs |
| 1214 | if not item.is_splitting_graph |
| 1215 | ] |
| 1216 | |
| 1217 | # Extract fake values from the graph to use them when needed. |
| 1218 | all_fake_values = [] |
| 1219 | for i in graph.graph.find_nodes(op="placeholder"): |
| 1220 | all_fake_values.append(i.meta["example_value"]) |
| 1221 | |
| 1222 | fake_args = [ |
| 1223 | all_fake_values[i] if isinstance(t, torch.Tensor) else t |
| 1224 | for i, t in enumerate(example_inputs) |
| 1225 | ] |
| 1226 | |
| 1227 | # propagate the split graph to the piecewise backend, |
| 1228 | # compile submodules with symbolic shapes, and compile all ranges |
| 1229 | # up front so that compilation is complete before the callable |
| 1230 | # is returned. |
| 1231 | PiecewiseCompileInterpreter( |
| 1232 | self.split_gm, submod_names_to_compile, self.vllm_config, self |
| 1233 | ).run(*fake_args) |
| 1234 | |
| 1235 | # All compilation is done. Save the cache. |
| 1236 | time_before_saving = time.perf_counter() |
| 1237 | self.compiler_manager.save_to_file() |
| 1238 | elapsed = time.perf_counter() - time_before_saving |
| 1239 | if elapsed > 1: |
| 1240 | logger.info_once( |
| 1241 | "Saved compiler manager cache in %.2f seconds.", |
| 1242 | elapsed, |
| 1243 | ) |
| 1244 | |
| 1245 | from torch._guards import detect_fake_mode |
| 1246 | |
| 1247 | fake_mode = detect_fake_mode() |
| 1248 | |
| 1249 | if ( |
| 1250 | self.compilation_config.dynamic_shapes_config.evaluate_guards |
| 1251 | and self.compilation_config.dynamic_shapes_config.type |
| 1252 | == DynamicShapesType.BACKED |
| 1253 | ): |
| 1254 | from torch.utils._sympy.value_ranges import ValueRanges |
| 1255 | |
| 1256 | # Drop counter-0/1 specializations guards; for backed dynamic shapes, |
| 1257 | # torch.compile will specialize for 0/1 inputs or otherwise guards that |
| 1258 | # shape is >= 2. This is because it's really hard not to hit a check |
| 1259 | # against 0/1. When we evaluate shape guards, we exclude checking those |
| 1260 | # guards (We would fail always otherwise). |
| 1261 | |
| 1262 | # We avoid that by updating the ranges of backed sizes when the min is |
| 1263 | # 2 for any, we assume it's 0. |
| 1264 | for s, r in fake_mode.shape_env.var_to_range.items(): |
| 1265 | if r.lower == 2: |
| 1266 | fake_mode.shape_env.var_to_range[s] = ValueRanges(0, r.upper) |
| 1267 | |
| 1268 | graph_path = os.path.join(local_cache_dir, "computation_graph.py") |
| 1269 | if not os.path.exists(graph_path): |
| 1270 | # code adapted from |
| 1271 | # https://github.com/thuml/depyf/blob/dab831108a752d1facc00acdd6d4243891845c37/depyf/explain/patched_lazy_format_graph_code.py#L30 |
| 1272 | # use `print_readable` because it can include submodules |
| 1273 | with dynamo_timed("vllm_print_readable"): |
| 1274 | src = ( |
| 1275 | "from __future__ import annotations\nimport torch\n" |
| 1276 | + self.split_gm.print_readable(print_output=False) |
| 1277 | ) |
| 1278 | src = src.replace("<lambda>", "GraphModule") |
| 1279 | with open(graph_path, "w") as f: |
| 1280 | f.write(src) |
| 1281 | |
| 1282 | logger.debug_once("Computation graph saved to %s", graph_path) |
| 1283 | |
| 1284 | self._called = True |
| 1285 | graph_to_serialize = ( |
| 1286 | original_split_gm if envs.VLLM_USE_MEGA_AOT_ARTIFACT else self.graph |
| 1287 | ) |
| 1288 | |
| 1289 | execution_code, submod_names, consts = generate_execution_code(self.split_gm) |
| 1290 | # Use getattr to get correct callables: __dict__ has PiecewiseBackend |
| 1291 | # instances (from PiecewiseCompileInterpreter), _modules has originals. |
| 1292 | # getattr checks __dict__ first, then falls back to _modules. |
| 1293 | submod_callables = { |
| 1294 | name: getattr(self.split_gm, name) |
| 1295 | for name, _ in self.split_gm.named_children() |
| 1296 | } |
| 1297 | runtime_callable = compile_execution_fn( |
| 1298 | execution_code, submod_callables, submod_names, consts |
| 1299 | ) |
| 1300 | |
| 1301 | if ( |
| 1302 | self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE |
| 1303 | or not self.compilation_config.cudagraph_copy_inputs |
| 1304 | ): |
| 1305 | return VllmSerializableFunction( |
| 1306 | graph_to_serialize, |
| 1307 | example_inputs, |
| 1308 | self.prefix, |
| 1309 | runtime_callable, |
| 1310 | is_encoder=self.is_encoder, |
| 1311 | vllm_backend=self, |
| 1312 | execution_code=execution_code, |
| 1313 | submod_names=submod_names, |
| 1314 | consts=consts, |
| 1315 | ) |
| 1316 | |
| 1317 | # index of tensors that have symbolic shapes (batch size) |
| 1318 | # for weights and static buffers, they will have concrete shapes. |
| 1319 | # symbolic shape only happens for input tensors. |
| 1320 | from torch.fx.experimental.symbolic_shapes import is_symbolic |
| 1321 | |
| 1322 | sym_tensor_indices = [ |
| 1323 | i |
| 1324 | for i, x in enumerate(fake_args) |
| 1325 | if isinstance(x, torch._subclasses.fake_tensor.FakeTensor) |
| 1326 | and any(is_symbolic(d) for d in x.size()) |
| 1327 | ] |
| 1328 | |
| 1329 | # compiler managed cudagraph input buffers |
| 1330 | # we assume the first run with symbolic shapes |
| 1331 | # has the maximum size among all the tensors |
| 1332 | copy_and_call = make_copy_and_call( |
| 1333 | sym_tensor_indices, |
| 1334 | [example_inputs[x].clone() for x in sym_tensor_indices], |
| 1335 | runtime_callable, |
| 1336 | ) |
| 1337 | |
| 1338 | return VllmSerializableFunction( |
| 1339 | graph_to_serialize, |
| 1340 | example_inputs, |
| 1341 | self.prefix, |
| 1342 | copy_and_call, |
| 1343 | is_encoder=self.is_encoder, |
| 1344 | vllm_backend=self, |
| 1345 | sym_tensor_indices=sym_tensor_indices, |
| 1346 | execution_code=execution_code, |
| 1347 | submod_names=submod_names, |
| 1348 | consts=consts, |
| 1349 | ) |