| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import contextlib |
| 5 | import hashlib |
| 6 | import inspect |
| 7 | import os |
| 8 | import sys |
| 9 | from collections.abc import Callable, Generator |
| 10 | from typing import TYPE_CHECKING, Any, TypeAlias, TypeVar, overload |
| 11 | from unittest.mock import patch |
| 12 | |
| 13 | import torch |
| 14 | import torch.nn as nn |
| 15 | from torch._dynamo.symbolic_convert import InliningInstructionTranslator |
| 16 | |
| 17 | import vllm.envs as envs |
| 18 | from vllm.compilation.counter import compilation_counter |
| 19 | from vllm.compilation.wrapper import TorchCompileWithNoGuardsWrapper |
| 20 | from vllm.config import ( |
| 21 | CompilationMode, |
| 22 | VllmConfig, |
| 23 | get_current_vllm_config, |
| 24 | set_current_vllm_config, |
| 25 | ) |
| 26 | from vllm.config.compilation import DynamicShapesType |
| 27 | from vllm.forward_context import get_forward_context, is_forward_context_available |
| 28 | from vllm.logger import init_logger |
| 29 | from vllm.sequence import IntermediateTensors |
| 30 | from vllm.utils.import_utils import resolve_obj_by_qualname |
| 31 | from vllm.utils.torch_utils import is_torch_equal_or_newer |
| 32 | |
| 33 | from .monitor import monitor_profiling_run, monitor_torch_compile |
| 34 | |
| 35 | DynamicArgDims: TypeAlias = dict[str, int | list[int] | dict[int, str]] |
| 36 | """Argument name -> the dynamic dimension(s) of that argument.""" |
| 37 | |
| 38 | # shape_id parameter was added to mark_unbacked in PyTorch 2.11.0 |
| 39 | _SUPPORTS_SHAPE_ID = is_torch_equal_or_newer("2.11.0") |
| 40 | |
| 41 | if TYPE_CHECKING: |
| 42 | # Only added on nightly/2.10 so wrap |
| 43 | try: |
| 44 | from torch._dynamo.package import SourceInfo |
| 45 | except ImportError: |
| 46 | # Fallback for old versions not supporting |
| 47 | SourceInfo = Any |
| 48 | |
| 49 | logger = init_logger(__name__) |
| 50 | |
| 51 | IGNORE_COMPILE_KEY = "_ignore_compile_vllm" |
| 52 | |
| 53 | _T = TypeVar("_T", bound=nn.Module) |
| 54 | |
| 55 | |
| 56 | def should_torch_compile_mm_encoder(vllm_config: VllmConfig) -> bool: |
| 57 | """Callable to be passed to `@support_torch_compile`'s `enable_if` argument.""" |
| 58 | return vllm_config.compilation_config.compile_mm_encoder |
| 59 | |
| 60 | |
| 61 | def ignore_torch_compile(cls: type[_T]) -> type[_T]: |
| 62 | """A decorator to ignore support_torch_compile decorator |
| 63 | on the class. This is useful when a parent class has |
| 64 | a support_torch_compile decorator, but we don't want to |
| 65 | compile the class `cls` that inherits the parent class. |
| 66 | This only ignores compiling the forward of the class the |
| 67 | decorator is applied to. |
| 68 | |
| 69 | If the parent has ignore_torch_compile but the child has |
| 70 | support_torch_compile, the child will still be compiled. |
| 71 | |
| 72 | If the class has one or more submodules |
| 73 | that have support_torch_compile decorator applied, compile will |
| 74 | not be ignored for those submodules. |
| 75 | """ |
| 76 | setattr(cls, IGNORE_COMPILE_KEY, True) |
| 77 | return cls |
| 78 | |
| 79 | |
| 80 | def _should_ignore_torch_compile(cls: type[_T]) -> bool: |
| 81 | """Check if the class should be ignored for torch.compile.""" |
| 82 | return getattr(cls, IGNORE_COMPILE_KEY, False) |
| 83 | |
| 84 | |
| 85 | @overload |
| 86 | def support_torch_compile( |
| 87 | *, |
| 88 | enable_if: Callable[[VllmConfig], bool] | None = None, |
| 89 | ) -> Callable[[type[_T]], type[_T]]: ... |
| 90 | |
| 91 | |
| 92 | @overload |
| 93 | def support_torch_compile( |
| 94 | *, |
| 95 | dynamic_arg_dims: DynamicArgDims | None, |
| 96 | ) -> Callable[[type[_T]], type[_T]]: ... |
| 97 | |
| 98 | |
| 99 | @overload |
| 100 | def support_torch_compile( |
| 101 | *, |
| 102 | mark_unbacked_dims: dict[str, int | list[int]] | None, |
| 103 | ) -> Callable[[type[_T]], type[_T]]: ... |
| 104 | |
| 105 | |
| 106 | @overload |
| 107 | def support_torch_compile( |
| 108 | *, |
| 109 | dynamic_arg_dims: DynamicArgDims | None = None, |
| 110 | mark_unbacked_dims: dict[str, int | list[int]] | None = None, |
| 111 | enable_if: Callable[[VllmConfig], bool] | None = None, |
| 112 | is_encoder: bool = False, |
| 113 | ) -> Callable[[type[_T]], type[_T]]: ... |
| 114 | |
| 115 | |
| 116 | @overload |
| 117 | def support_torch_compile(cls: type[_T]) -> type[_T]: ... |
| 118 | |
| 119 | |
| 120 | def support_torch_compile( |
| 121 | cls: type[_T] | None = None, |
| 122 | *, |
| 123 | dynamic_arg_dims: DynamicArgDims | None = None, |
| 124 | mark_unbacked_dims: dict[str, int | list[int]] | None = None, |
| 125 | enable_if: Callable[[VllmConfig], bool] | None = None, |
| 126 | is_encoder: bool = False, |
| 127 | ) -> Callable[[type[_T]], type[_T]] | type[_T]: |
| 128 | """A decorator to add support for compiling the forward method of a class. |
| 129 | |
| 130 | Usage 1: use directly as a decorator without arguments: |
| 131 | |
| 132 | ```python |
| 133 | @support_torch_compile |
| 134 | class MyModel(nn.Module): |
| 135 | def forward(self, x: torch.Tensor, y: Optional[torch.Tensor]): ... |
| 136 | ``` |
| 137 | |
| 138 | Usage 2: use as a decorator with arguments: |
| 139 | |
| 140 | ```python |
| 141 | @support_torch_compile(dynamic_arg_dims={"x": 0, "y": 0}) |
| 142 | class MyModel(nn.Module): |
| 143 | def forward(self, x: torch.Tensor, y: Optional[torch.Tensor]): ... |
| 144 | ``` |
| 145 | |
| 146 | `dynamic_arg_dims` is a dictionary that maps argument names to the dynamic |
| 147 | dimensions of the argument. The value can be: |
| 148 | - int: a single dimension index (e.g., 0) |
| 149 | - list[int]: multiple dimension indices (e.g., [0, 1]) |
| 150 | - dict[int, str]: dimension to shape_id mapping for shape relations |
| 151 | (e.g., {0: "b"}). Dimensions with the same shape_id share the same |
| 152 | unbacked symbol. |
| 153 | |
| 154 | if `dynamic_arg_dims` is `None`, it is inferred from the type annotation |
| 155 | of the `forward` method, based on the following default rules: |
| 156 | |
| 157 | - if the argument is annotated as `torch.Tensor` or |
| 158 | `Optional[torch.Tensor]`, the first dimension will be |
| 159 | marked as dynamic. |
| 160 | - if the argument is annotated as `IntermediateTensors`, the first |
| 161 | dimension of all the tensors in the intermediate tensors |
| 162 | will be marked as dynamic. |
| 163 | |
| 164 | During runtime, when we actually mark dimensions of tensors, |
| 165 | it depends on the value of arguments: |
| 166 | |
| 167 | - if it is a single integer (can be negative), the corresponding dimension |
| 168 | of the argument will be marked as dynamic. |
| 169 | - if it is `None`, ignored. |
| 170 | - if it is `IntermediateTensors`, all the tensors in the intermediate |
| 171 | tensors will be marked as dynamic. |
| 172 | - otherwise, it will raise an error. |
| 173 | |
| 174 | NOTE: if an argument is `None`, it should always be passed as `None` during |
| 175 | the lifetime of the model, otherwise, it cannot be captured as a single |
| 176 | computation graph. |
| 177 | |
| 178 | `enable_if` is a function that takes a `VllmConfig` object as input and |
| 179 | returns a boolean value indicating whether to compile the model or not. |
| 180 | This is useful if you want to compile the model only when certain |
| 181 | conditions are met. |
| 182 | |
| 183 | `mark_unbacked_dims` is a dictionary that maps argument names with a dynamic |
| 184 | dim to be decorated with `mark_unbacked`. This is useful if we would like to |
| 185 | enforce that dynamo does not specialize on 0/1 values in the case of dummy input |
| 186 | such as for vision model compilation |
| 187 | |
| 188 | `is_encoder` marks this module as a portion of an multimodal encoder. |
| 189 | When True, the compile range upper bound is set to MAX_INT32 instead of |
| 190 | max_num_batched_tokens, since encoder input shapes are unpredictable. |
| 191 | This is typically used for vision encoder sub-modules in multimodal models. |
| 192 | |
| 193 | `shape_invariants` is a function that gets compiled right before forward. |
| 194 | The function should have the torch._check calls that are needed to set |
| 195 | the relationships between different input sizes. For example: |
| 196 | torch._check(input_ids.size()[0] == inputs_embeds.size()[0]) |
| 197 | This enforces constraints on the symbolic shapes without hardcoding |
| 198 | specific values. It is needed for some models to avoid data dependent |
| 199 | errors and maximize perf when unbacked shapes are used. |
| 200 | """ |
| 201 | |
| 202 | def cls_decorator_helper(cls: type[_T]) -> type[_T]: |
| 203 | # helper to pass `dynamic_arg_dims` to `_support_torch_compile` |
| 204 | # to avoid too much indentation for `_support_torch_compile` |
| 205 | if not hasattr(cls, "forward"): |
| 206 | raise TypeError("decorated class should have a forward method.") |
| 207 | sig = inspect.signature(cls.forward) |
| 208 | inferred_dynamic_arg_dims = dynamic_arg_dims |
| 209 | if inferred_dynamic_arg_dims is None: |
| 210 | inferred_dynamic_arg_dims = {} |
| 211 | for k, v in sig.parameters.items(): |
| 212 | if v.annotation in [ |
| 213 | torch.Tensor, |
| 214 | torch.Tensor | None, |
| 215 | torch.FloatTensor, |
| 216 | torch.FloatTensor | None, |
| 217 | IntermediateTensors, |
| 218 | IntermediateTensors | None, |
| 219 | ]: |
| 220 | inferred_dynamic_arg_dims[k] = 0 |
| 221 | |
| 222 | logger.debug( |
| 223 | ("Inferred dynamic dimensions for forward method of %s: %s"), |
| 224 | cls, |
| 225 | list(inferred_dynamic_arg_dims.keys()), |
| 226 | ) |
| 227 | |
| 228 | if len(inferred_dynamic_arg_dims) == 0: |
| 229 | raise ValueError( |
| 230 | "No dynamic dimensions found in the forward method of " |
| 231 | f"{cls}. Please provide dynamic_arg_dims explicitly." |
| 232 | ) |
| 233 | |
| 234 | for k in inferred_dynamic_arg_dims: |
| 235 | if k not in sig.parameters: |
| 236 | raise ValueError( |
| 237 | f"Argument {k} not found in the forward method of {cls}" |
| 238 | ) |
| 239 | |
| 240 | return _support_torch_compile( |
| 241 | cls, |
| 242 | inferred_dynamic_arg_dims, |
| 243 | mark_unbacked_dims, |
| 244 | enable_if, |
| 245 | is_encoder, |
| 246 | ) |
| 247 | |
| 248 | if cls is not None: |
| 249 | # use `support_torch_compile` as a decorator without arguments |
| 250 | assert isinstance(cls, type) |
| 251 | return cls_decorator_helper(cls) |
| 252 | |
| 253 | return cls_decorator_helper |
| 254 | |
| 255 | |
| 256 | def _model_hash_key(fn: Callable[..., Any]) -> str: |
| 257 | import vllm |
| 258 | |
| 259 | sha256_hash = hashlib.sha256() |
| 260 | sha256_hash.update(vllm.__version__.encode()) |
| 261 | sha256_hash.update(fn.__qualname__.encode()) |
| 262 | sha256_hash.update(str(fn.__code__.co_firstlineno).encode()) |
| 263 | return sha256_hash.hexdigest() |
| 264 | |
| 265 | |
| 266 | def _verify_source_unchanged( |
| 267 | source_info: "SourceInfo", vllm_config: VllmConfig |
| 268 | ) -> None: |
| 269 | from .caching import _compute_code_hash, _compute_code_hash_with_content |
| 270 | |
| 271 | file_contents = {} |
| 272 | for source in source_info.inlined_sources: |
| 273 | module = sys.modules[source.module] |
| 274 | file = inspect.getfile(module) |
| 275 | vllm_config.compilation_config.traced_files.add(file) |
| 276 | file_contents[file] = source.content |
| 277 | expected_checksum = _compute_code_hash_with_content(file_contents) |
| 278 | actual_checksum = _compute_code_hash(set(file_contents.keys())) |
| 279 | if expected_checksum != actual_checksum: |
| 280 | raise RuntimeError( |
| 281 | "Source code has changed since the last compilation. Recompiling the model." |
| 282 | ) |
| 283 | |
| 284 | |
| 285 | def _try_load_aot_compiled_fn( |
| 286 | model: Any, |
| 287 | aot_compilation_path: str, |
| 288 | ) -> Any | None: |
| 289 | """Try to load an AOT-compiled function from disk. |
| 290 | |
| 291 | Returns the loaded callable on success, or None on failure. |
| 292 | Re-raises on failure when ``VLLM_FORCE_AOT_LOAD`` is set. |
| 293 | """ |
| 294 | try: |
| 295 | with monitor_torch_compile( |
| 296 | model.vllm_config, |
| 297 | is_encoder=model._is_encoder, |
| 298 | tag=getattr(model, "_compile_tag", ""), |
| 299 | announce_start=False, |
| 300 | ): |
| 301 | with ( |
| 302 | set_current_vllm_config(model.vllm_config), |
| 303 | open(aot_compilation_path, "rb") as f, |
| 304 | ): |
| 305 | loaded_fn = torch.compiler.load_compiled_function( |
| 306 | f, f_globals=model.forward.__globals__ |
| 307 | ) |
| 308 | _verify_source_unchanged(loaded_fn.source_info(), model.vllm_config) |
| 309 | ds_config = model.compilation_config.dynamic_shapes_config |
| 310 | if not ds_config.evaluate_guards: |
| 311 | loaded_fn.disable_guard_check() |
| 312 | # Eagerly load compiled artifacts now that traced_files |
| 313 | # is populated by _verify_source_unchanged. |
| 314 | with maybe_use_cudagraph_partition_wrapper(model.vllm_config): |
| 315 | loaded_fn._artifacts.compiled_fn.finalize_loading(model.vllm_config) |
| 316 | compilation_counter.num_aot_artifacts_loaded += 1 |
| 317 | tag = getattr(model, "_compile_tag", "") |
| 318 | log_prefix = f"[{tag}] " if tag else "" |
| 319 | logger.info( |
| 320 | "%sDirectly load AOT compilation from path %s", |
| 321 | log_prefix, |
| 322 | aot_compilation_path, |
| 323 | ) |
| 324 | return loaded_fn |
| 325 | except Exception as e: |
| 326 | if os.path.exists(aot_compilation_path): |
| 327 | if isinstance(e, EOFError): |
| 328 | message = "Compile cache file corrupted." |
| 329 | else: |
| 330 | message = str(e) |
| 331 | tag = getattr(model, "_compile_tag", "") |
| 332 | log_prefix = f"[{tag}] " if tag else "" |
| 333 | logger.warning( |
| 334 | "%sCompiling model again due to a load failure from %s, reason: %s", |
| 335 | log_prefix, |
| 336 | aot_compilation_path, |
| 337 | message, |
| 338 | ) |
| 339 | if envs.VLLM_FORCE_AOT_LOAD: |
| 340 | raise e |
| 341 | return None |
| 342 | |
| 343 | |
| 344 | def _support_torch_compile( |
| 345 | cls: type[_T], |
| 346 | dynamic_arg_dims: DynamicArgDims, |
| 347 | mark_unbacked_dims: dict[str, int | list[int]] | None = None, |
| 348 | enable_if: Callable[[VllmConfig], bool] | None = None, |
| 349 | is_encoder: bool = False, |
| 350 | ) -> type[_T]: |
| 351 | """Internal implementation of support_torch_compile decorator.""" |
| 352 | if TorchCompileWithNoGuardsWrapper in cls.__bases__: |
| 353 | # support decorating multiple times |
| 354 | return cls |
| 355 | |
| 356 | # take care of method resolution order |
| 357 | # make sure super().__init__ is called on the base class |
| 358 | # other than TorchCompileWithNoGuardsWrapper |
| 359 | cls.__bases__ = cls.__bases__ + (TorchCompileWithNoGuardsWrapper,) |
| 360 | |
| 361 | old_init = cls.__init__ |
| 362 | |
| 363 | setattr(cls, IGNORE_COMPILE_KEY, False) |
| 364 | |
| 365 | def __init__( |
| 366 | self: _T, |
| 367 | *args, |
| 368 | vllm_config: VllmConfig | None = None, |
| 369 | prefix: str = "", |
| 370 | **kwargs: Any, |
| 371 | ) -> None: |
| 372 | if vllm_config is None: |
| 373 | vllm_config = get_current_vllm_config() |
| 374 | |
| 375 | # NOTE: to support multimodal models (such as encoder), |
| 376 | # we may not have vllm_config so we may need to patch it |
| 377 | sig = inspect.signature(old_init) |
| 378 | # Check that any positional arguments match the old_init method signature |
| 379 | annotations = [p.annotation for p in sig.parameters.values()] |
| 380 | for arg, annotation in zip(args, annotations): |
| 381 | if annotation is inspect._empty: |
| 382 | continue |
| 383 | if not isinstance(arg, annotation): |
| 384 | init = f"'{type(self).__name__}.__init__'" |
| 385 | arg_type = f"'{type(arg).__name__}'" |
| 386 | raise TypeError( |
| 387 | f"{init} received a positional argument of type {arg_type}, " |
| 388 | "but no parameter of that type was found in the method signature. " |
| 389 | f"Please either annotate {init} or pass it as a keyword argument." |
| 390 | ) |
| 391 | if "vllm_config" in sig.parameters: |
| 392 | kwargs["vllm_config"] = vllm_config |
| 393 | if "prefix" in sig.parameters: |
| 394 | kwargs["prefix"] = prefix |
| 395 | old_init(self, *args, **kwargs) |
| 396 | |
| 397 | self.vllm_config = vllm_config |
| 398 | self.compilation_config = self.vllm_config.compilation_config |
| 399 | enable_compile = enable_if is None or enable_if(vllm_config) |
| 400 | # for CompilationMode.STOCK_TORCH_COMPILE , the upper level model runner |
| 401 | # will handle the compilation, so we don't need to do anything here. |
| 402 | self.do_not_compile = ( |
| 403 | self.compilation_config.mode |
| 404 | in [CompilationMode.NONE, CompilationMode.STOCK_TORCH_COMPILE] |
| 405 | or _should_ignore_torch_compile(self.__class__) |
| 406 | or not enable_compile |
| 407 | ) |
| 408 | if self.do_not_compile: |
| 409 | return |
| 410 | |
| 411 | self._dynamic_arg_dims = dynamic_arg_dims |
| 412 | |
| 413 | self.was_aot_compile_fn_loaded_from_disk = False |
| 414 | compilation_counter.num_models_seen += 1 |
| 415 | self.compiled = False |
| 416 | |
| 417 | # Capture the active model tag (e.g. "backbone", "eagle_head") so |
| 418 | # later compilation/warmup log lines can identify which component |
| 419 | # they belong to. Mirror VllmBackend's precedence rule |
| 420 | # (`compile_prefix or model_tag`) so log tags line up with the |
| 421 | # cache-dir tag: encoders use the class name, everything else |
| 422 | # falls through to `model_tag`. |
| 423 | from vllm.compilation.backends import model_tag as _current_model_tag |
| 424 | |
| 425 | self._compile_tag = (cls.__name__ if is_encoder else "") or _current_model_tag |
| 426 | |
| 427 | # Handled by monkeypatching `TorchCompileWithNoGuardsWrapper` into base class |
| 428 | TorchCompileWithNoGuardsWrapper.__init__( |
| 429 | self, |
| 430 | compile_prefix=cls.__name__ if is_encoder else "", |
| 431 | is_encoder=is_encoder, |
| 432 | ) |
| 433 | |
| 434 | cls.__init__ = __init__ |
| 435 | |
| 436 | def _mark_dynamic_inputs( |
| 437 | mod: type[_T], ds_type: DynamicShapesType, *args: Any, **kwargs: Any |
| 438 | ) -> None: |
| 439 | def mark_dynamic( |
| 440 | arg: torch.Tensor, dim_shape_pairs: list[tuple[int, str | None]] |
| 441 | ) -> None: |
| 442 | if ds_type == DynamicShapesType.UNBACKED: |
| 443 | if is_torch_equal_or_newer("2.10.0"): |
| 444 | for dim, shape_id in dim_shape_pairs: |
| 445 | if shape_id is not None: |
| 446 | if not _SUPPORTS_SHAPE_ID: |
| 447 | raise RuntimeError( |
| 448 | f"shape_id='{shape_id}' requires PyTorch >= 2.11.0" |
| 449 | ) |
| 450 | torch._dynamo.decorators.mark_unbacked( |
| 451 | arg, |
| 452 | dim, |
| 453 | hint_override=arg.size()[dim], |
| 454 | shape_id=shape_id, |
| 455 | ) |
| 456 | else: |
| 457 | torch._dynamo.decorators.mark_unbacked( |
| 458 | arg, |
| 459 | dim, |
| 460 | hint_override=arg.size()[dim], |
| 461 | ) |
| 462 | else: |
| 463 | # For older versions, we can't use hint_override or shape_id |
| 464 | dims = [dim for dim, _ in dim_shape_pairs] |
| 465 | torch._dynamo.decorators.mark_unbacked(arg, dims) |
| 466 | else: |
| 467 | dims = [dim for dim, _ in dim_shape_pairs] |
| 468 | torch._dynamo.mark_dynamic(arg, dims) |
| 469 | |
| 470 | sig = inspect.signature(mod.__class__.forward) # type: ignore[attr-defined] |
| 471 | bound_args = sig.bind(mod, *args, **kwargs) |
| 472 | bound_args.apply_defaults() |
| 473 | |
| 474 | # Normalize dynamic_arg_dims to dict[str, dict[int, str | None]] |
| 475 | normalized_dims: dict[str, dict[int, str | None]] = {} |
| 476 | for k, v in dynamic_arg_dims.items(): |
| 477 | if isinstance(v, dict): |
| 478 | normalized_dims[k] = {dim: shape_id for dim, shape_id in v.items()} |
| 479 | elif isinstance(v, int): |
| 480 | normalized_dims[k] = {v: None} |
| 481 | else: |
| 482 | normalized_dims[k] = {d: None for d in v} |
| 483 | |
| 484 | for k, dim_to_shape_id in normalized_dims.items(): |
| 485 | arg = bound_args.arguments.get(k) |
| 486 | |
| 487 | if arg is not None: |
| 488 | dims = list(dim_to_shape_id.keys()) |
| 489 | |
| 490 | if isinstance(arg, torch.Tensor): |
| 491 | dim_shape_pairs = [ |
| 492 | (arg.ndim + d if d < 0 else d, dim_to_shape_id.get(d)) |
| 493 | for d in dims |
| 494 | ] |
| 495 | mark_dynamic(arg, dim_shape_pairs) |
| 496 | elif isinstance(arg, IntermediateTensors): |
| 497 | for tensor in arg.tensors.values(): |
| 498 | dim_shape_pairs = [ |
| 499 | (tensor.ndim + d if d < 0 else d, dim_to_shape_id.get(d)) |
| 500 | for d in dims |
| 501 | ] |
| 502 | mark_dynamic(tensor, dim_shape_pairs) |
| 503 | else: |
| 504 | raise ValueError( |
| 505 | f"Unsupported dynamic dimensions {dims} " |
| 506 | f"for argument {k} with type {type(arg)}." |
| 507 | ) |
| 508 | |
| 509 | if mark_unbacked_dims: |
| 510 | for k, dims_val in mark_unbacked_dims.items(): |
| 511 | arg = bound_args.arguments.get(k) |
| 512 | if arg is not None: |
| 513 | dims = [dims_val] if isinstance(dims_val, int) else list(dims_val) |
| 514 | if isinstance(arg, torch.Tensor): |
| 515 | dims = [arg.ndim + d if d < 0 else d for d in dims] |
| 516 | if is_torch_equal_or_newer("2.10.0"): |
| 517 | for dim in dims: |
| 518 | torch._dynamo.decorators.mark_unbacked( |
| 519 | arg, dim, hint_override=arg.size()[dim] |
| 520 | ) |
| 521 | else: |
| 522 | torch._dynamo.decorators.mark_unbacked(arg, dims) |
| 523 | |
| 524 | def __call__(self: type[_T], *args: Any, **kwargs: Any) -> Any: |
| 525 | # torch.compiler.is_compiling() means we are inside the compilation |
| 526 | # e.g. TPU has the compilation logic in model runner, so we don't |
| 527 | # need to compile the model inside. |
| 528 | if self.do_not_compile or torch.compiler.is_compiling(): |
| 529 | return self.forward(*args, **kwargs) |
| 530 | |
| 531 | # If skip_compiled is set, bypass compiled model call. This is used e.g. for |
| 532 | # enc-dec models where tensor shapes/types vary across invocations, preventing |
| 533 | # the capture of a single computational graph. |
| 534 | if is_forward_context_available() and get_forward_context().skip_compiled: |
| 535 | return self.forward(*args, **kwargs) |
| 536 | |
| 537 | # if aot_compiled_fn is set, call it with partition wrapper context. |
| 538 | # The partition wrapper must be active at runtime for CUDA graph |
| 539 | # capture to work correctly with inductor graph partitioning. |
| 540 | if getattr(self, "aot_compiled_fn", None) is not None: |
| 541 | with maybe_use_cudagraph_partition_wrapper(self.vllm_config): |
| 542 | return self.aot_compiled_fn(self, *args, **kwargs) |
| 543 | |
| 544 | ds_type = self.compilation_config.dynamic_shapes_config.type |
| 545 | cache_dir = None |
| 546 | aot_compilation_path = None |
| 547 | if envs.VLLM_USE_AOT_COMPILE: |
| 548 | """ |
| 549 | When using torch.compile in AOT mode, we store the cache artifacts |
| 550 | under VLLM_CACHE_ROOT/torch_compile_cache/torch_aot_compile/{hash} |
| 551 | The {hash} contains all of the factors except for the source files |
| 552 | being traced through, because we don't actually know which source |
| 553 | files to check at this point (before dynamo runs). |
| 554 | On loading we will actually look at the source files being traced |
| 555 | through. If any source file have changed (compared with the |
| 556 | serialized backend artifacts), then we need to generate a new AOT |
| 557 | compile artifact from scratch. |
| 558 | """ |
| 559 | from .caching import aot_compile_hash_factors |
| 560 | |
| 561 | factors: list[str] = aot_compile_hash_factors(self.vllm_config) |
| 562 | |
| 563 | factors.append(_model_hash_key(self.forward)) |
| 564 | hash_key = hashlib.sha256(str(factors).encode()).hexdigest() |
| 565 | cache_dir = os.path.join( |
| 566 | envs.VLLM_CACHE_ROOT, |
| 567 | "torch_compile_cache", |
| 568 | "torch_aot_compile", |
| 569 | hash_key, |
| 570 | ) |
| 571 | |
| 572 | # Hash-level dir; shared across ranks on the same node. |
| 573 | self.compilation_config.local_cache_dir = cache_dir |
| 574 | inductor_cache = os.path.join(cache_dir, "inductor_cache") |
| 575 | os.makedirs(inductor_cache, exist_ok=True) |
| 576 | # Process-wide: post-load execution, CUDA-graph capture, and later |
| 577 | # autotune/recompile all need to write under {hash}/inductor_cache/. |
| 578 | # Unconditional because torch's cache_dir() may have pre-filled the |
| 579 | # /tmp default during import, making setdefault a no-op. |
| 580 | os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache |
| 581 | |
| 582 | rank = self.vllm_config.parallel_config.rank |
| 583 | dp_rank = self.vllm_config.parallel_config.data_parallel_index |
| 584 | cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}") |
| 585 | aot_compilation_path = os.path.join(cache_dir, "model") |
| 586 | if not envs.VLLM_DISABLE_COMPILE_CACHE: |
| 587 | loaded_fn = _try_load_aot_compiled_fn(self, aot_compilation_path) |
| 588 | if loaded_fn is not None: |
| 589 | self.aot_compiled_fn = loaded_fn |
| 590 | self.was_aot_compile_fn_loaded_from_disk = True |
| 591 | with ( |
| 592 | monitor_profiling_run(tag=self._compile_tag), |
| 593 | maybe_use_cudagraph_partition_wrapper(self.vllm_config), |
| 594 | ): |
| 595 | output = self.aot_compiled_fn(self, *args, **kwargs) |
| 596 | return output |
| 597 | |
| 598 | if self.compiled: |
| 599 | assert ( |
| 600 | not envs.VLLM_USE_AOT_COMPILE |
| 601 | or self.vllm_config.compilation_config.backend == "eager" |
| 602 | ) |
| 603 | return TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs) # type: ignore[arg-type] |
| 604 | |
| 605 | # This is the path for the first compilation. |
| 606 | # the first compilation needs to have dynamic shapes marked |
| 607 | _mark_dynamic_inputs( |
| 608 | self, |
| 609 | ds_type, |
| 610 | *args, |
| 611 | **kwargs, |
| 612 | ) |
| 613 | |
| 614 | # Tensors that enter the graph through the forward context rather than |
| 615 | # model args (e.g. the MoE padding mask) need their batch dim marked |
| 616 | # dynamic too, otherwise AOT-compiled artifacts bake a static guard |
| 617 | # for the compile-time batch size. |
| 618 | if is_forward_context_available(): |
| 619 | is_padding = get_forward_context().is_padding |
| 620 | if is_padding is not None: |
| 621 | torch._dynamo.mark_dynamic(is_padding, 0) |
| 622 | |
| 623 | original_code_object = self.original_code_object() |
| 624 | logger.debug("Start compiling function %s", original_code_object) |
| 625 | |
| 626 | # we do not want tp delete the original code object entries since |
| 627 | # we depend on them now to look up cached compiled functions. |
| 628 | # torch._dynamo.eval_frame.remove_from_cache(original_code_object) |
| 629 | |
| 630 | # collect all relevant files traced by Dynamo, |
| 631 | # so that the compilation cache can trigger re-compilation |
| 632 | # properly when any of these files change. |
| 633 | |
| 634 | # 1. the file containing the top-level forward function |
| 635 | self.compilation_config.traced_files.add(original_code_object.co_filename) |
| 636 | |
| 637 | # 2. every time Dynamo sees a function call, it will inline |
| 638 | # the function by calling InliningInstructionTranslator.inline_call_ |
| 639 | # we hijack this function to know all the functions called |
| 640 | # during Dynamo tracing, and their corresponding files |
| 641 | inline_call = InliningInstructionTranslator.inline_call_ |
| 642 | |
| 643 | def patched_inline_call(self_: Any) -> Any: |
| 644 | code = self_.f_code |
| 645 | self.compilation_config.traced_files.add(code.co_filename) |
| 646 | return inline_call(self_) |
| 647 | |
| 648 | # Disable the C++ compilation of symbolic shape guards. C++-fication |
| 649 | # of symbolic shape guards can improve guard overhead. But, since |
| 650 | # vllm skip guards anyways, setting this flag to False can improve |
| 651 | # compile time. |
| 652 | dynamo_config_patches = {} |
| 653 | try: |
| 654 | _ = torch._dynamo.config.enable_cpp_symbolic_shape_guards |
| 655 | dynamo_config_patches["enable_cpp_symbolic_shape_guards"] = False |
| 656 | except AttributeError: |
| 657 | # Note: this config is not available in torch 2.6, we can skip |
| 658 | # if the config doesn't exist |
| 659 | logger.debug("enable_cpp_symbolic_shape_guards config not available") |
| 660 | |
| 661 | # Prepare backed_size_oblivious config patch if needed |
| 662 | fx_config_patches = {} |
| 663 | if ds_type == DynamicShapesType.BACKED_SIZE_OBLIVIOUS: |
| 664 | fx_config_patches["backed_size_oblivious"] = True |
| 665 | |
| 666 | # Prepare inductor config patches |
| 667 | # assume_32bit_indexing is only available in torch 2.10.0+ |
| 668 | inductor_config_patches = {} |
| 669 | if is_torch_equal_or_newer("2.10.0"): |
| 670 | inductor_config_patches["assume_32bit_indexing"] = ( |
| 671 | self.compilation_config.dynamic_shapes_config.assume_32_bit_indexing |
| 672 | ) |
| 673 | |
| 674 | with ( |
| 675 | patch.object( |
| 676 | InliningInstructionTranslator, "inline_call_", patched_inline_call |
| 677 | ), |
| 678 | torch._dynamo.config.patch(**dynamo_config_patches), |
| 679 | maybe_use_cudagraph_partition_wrapper(self.vllm_config), |
| 680 | torch.fx.experimental._config.patch(**fx_config_patches), |
| 681 | torch._inductor.config.patch(**inductor_config_patches), |
| 682 | ): |
| 683 | use_aot_compile = envs.VLLM_USE_AOT_COMPILE |
| 684 | if self.vllm_config.compilation_config.backend == "eager": |
| 685 | logger.warning("Detected eager backend, disabling AOT compile.") |
| 686 | use_aot_compile = False |
| 687 | if use_aot_compile: |
| 688 | # store the path for saving after warmup |
| 689 | self._aot_compilation_path = aot_compilation_path |
| 690 | self._aot_cache_dir = cache_dir |
| 691 | with monitor_torch_compile( |
| 692 | self.vllm_config, |
| 693 | is_encoder=self._is_encoder, |
| 694 | tag=self._compile_tag, |
| 695 | ): |
| 696 | self.aot_compiled_fn = self.aot_compile(*args, **kwargs) |
| 697 | compilation_counter.num_aot_compiles += 1 |
| 698 | # All compilation is done at this point, save the |
| 699 | # AOT artifact. |
| 700 | self.save_aot_compiled_function() |
| 701 | |
| 702 | with monitor_profiling_run(tag=self._compile_tag): |
| 703 | output = self.aot_compiled_fn(self, *args, **kwargs) |
| 704 | else: |
| 705 | with monitor_torch_compile( |
| 706 | self.vllm_config, |
| 707 | "%storch.compile and initial profiling/warmup " |
| 708 | "run together took %.2f s in total", |
| 709 | is_encoder=self._is_encoder, |
| 710 | tag=self._compile_tag, |
| 711 | ): |
| 712 | output = TorchCompileWithNoGuardsWrapper.__call__( |
| 713 | self, # type: ignore[arg-type] |
| 714 | *args, |
| 715 | **kwargs, |
| 716 | ) |
| 717 | |
| 718 | self.compiled = True |
| 719 | return output |
| 720 | |
| 721 | # triggers VllmSerializableFunction.serialize() |
| 722 | def save_aot_compiled_function(self: type[_T]) -> None: |
| 723 | if envs.VLLM_DISABLE_COMPILE_CACHE: |
| 724 | return |
| 725 | |
| 726 | if self.was_aot_compile_fn_loaded_from_disk: |
| 727 | logger.debug("AOT compiled function was loaded from cache, skipping save") |
| 728 | return |
| 729 | |
| 730 | assert ( |
| 731 | self.aot_compiled_fn and self._aot_compilation_path and self._aot_cache_dir |
| 732 | ) |
| 733 | |
| 734 | log_prefix = f"[{self._compile_tag}] " if self._compile_tag else "" |
| 735 | try: |
| 736 | os.makedirs(self._aot_cache_dir, exist_ok=True) |
| 737 | # File saving should be atomic, so we will save to a temporary location |
| 738 | # first. Should be upstreamed to PyTorch 2.12 as well. |
| 739 | tmp_file = f"{self._aot_compilation_path}.{os.getpid()}.tmp" |
| 740 | self.aot_compiled_fn.save_compiled_function(tmp_file) |
| 741 | os.replace(tmp_file, self._aot_compilation_path) |
| 742 | compilation_counter.num_aot_artifacts_saved += 1 |
| 743 | logger.info( |
| 744 | "%ssaved AOT compiled function to %s", |
| 745 | log_prefix, |
| 746 | self._aot_compilation_path, |
| 747 | ) |
| 748 | except Exception as e: |
| 749 | logger.warning( |
| 750 | "%sunable to save AOT compiled function to %s: %s", |
| 751 | log_prefix, |
| 752 | self._aot_compilation_path, |
| 753 | e, |
| 754 | ) |
| 755 | |
| 756 | cls.__call__ = __call__ |
| 757 | cls.save_aot_compiled_function = save_aot_compiled_function |
| 758 | return cls |
| 759 | |
| 760 | |
| 761 | @contextlib.contextmanager |
| 762 | def maybe_use_cudagraph_partition_wrapper( |
| 763 | vllm_config: VllmConfig, |
| 764 | ) -> Generator[None, None, None]: |
| 765 | """Context manager to set/unset customized cudagraph partition wrappers. |
| 766 | |
| 767 | If we're using Inductor-based graph partitioning, we currently have the |
| 768 | whole `fx.Graph` before Inductor lowering and the piecewise |
| 769 | splitting happens after all graph passes and fusions. Here, we add |
| 770 | a custom hook for Inductor to wrap each partition with our static |
| 771 | graph wrapper class to maintain more control over static graph |
| 772 | capture and replay. |
| 773 | """ |
| 774 | from vllm.config import CUDAGraphMode |
| 775 | |
| 776 | compilation_config = vllm_config.compilation_config |
| 777 | if ( |
| 778 | compilation_config.cudagraph_mode.has_piecewise_cudagraphs() |
| 779 | and compilation_config.use_inductor_graph_partition |
| 780 | ): |
| 781 | from torch._inductor.utils import CUDAGraphWrapperMetadata |
| 782 | |
| 783 | from vllm.compilation.cuda_graph import CUDAGraphOptions |
| 784 | from vllm.platforms import current_platform |
| 785 | |
| 786 | static_graph_wrapper_class = resolve_obj_by_qualname( |
| 787 | current_platform.get_static_graph_wrapper_cls() |
| 788 | ) |
| 789 | |
| 790 | def customized_cudagraph_wrapper( |
| 791 | f: Callable[..., Any], metadata: CUDAGraphWrapperMetadata |
| 792 | ) -> Any: |
| 793 | partition_id = metadata.partition_index |
| 794 | num_partitions = metadata.num_partitions |
| 795 | return static_graph_wrapper_class( |
| 796 | runnable=f, |
| 797 | vllm_config=vllm_config, |
| 798 | runtime_mode=CUDAGraphMode.PIECEWISE, |
| 799 | cudagraph_options=CUDAGraphOptions( |
| 800 | debug_log_enable=partition_id == 0, |
| 801 | gc_disable=partition_id != 0, |
| 802 | weak_ref_output=partition_id == num_partitions - 1, |
| 803 | ), |
| 804 | ) |
| 805 | |
| 806 | torch._inductor.utils.set_customized_partition_wrappers( |
| 807 | customized_cudagraph_wrapper |
| 808 | ) |
| 809 | |
| 810 | yield |
| 811 | |
| 812 | if ( |
| 813 | compilation_config.cudagraph_mode.has_piecewise_cudagraphs() |
| 814 | and compilation_config.use_inductor_graph_partition |
| 815 | ): |
| 816 | torch._inductor.utils.set_customized_partition_wrappers(None) |