| 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 pickle |
| 9 | from collections.abc import Callable, Sequence |
| 10 | from typing import Any, Literal |
| 11 | from unittest.mock import patch |
| 12 | |
| 13 | import torch |
| 14 | from torch._subclasses import FakeTensorMode |
| 15 | from torch.fx._graph_pickler import GraphPickler, Options |
| 16 | from torch.utils import _pytree as pytree |
| 17 | |
| 18 | import vllm.envs as envs |
| 19 | from vllm.compilation.codegen import compile_execution_fn |
| 20 | from vllm.compilation.compiler_interface import get_inductor_factors |
| 21 | from vllm.compilation.counter import compilation_counter |
| 22 | from vllm.config import VllmConfig, get_current_vllm_config |
| 23 | from vllm.config.utils import hash_factors |
| 24 | from vllm.logger import init_logger |
| 25 | from vllm.platforms import current_platform |
| 26 | from vllm.utils.hashing import safe_hash |
| 27 | |
| 28 | try: |
| 29 | from torch._dynamo.aot_compile import SerializableCallable |
| 30 | except ImportError: |
| 31 | SerializableCallable = object |
| 32 | |
| 33 | assert isinstance(SerializableCallable, type) |
| 34 | |
| 35 | logger = init_logger(__name__) |
| 36 | |
| 37 | |
| 38 | class StandaloneCompiledArtifacts: |
| 39 | """Storage for standalone compiled artifacts with content-based deduplication. |
| 40 | |
| 41 | Deduplication works via a two-level indirection: |
| 42 | 1. `submodule_bytes` maps "{submod_name}_{shape}" -> SHA256 hash |
| 43 | 2. `submodule_bytes_store` maps SHA256 hash -> actual bytes |
| 44 | |
| 45 | When inserting, we compute the SHA256 hash of the bytes. If the hash |
| 46 | already exists in `submodule_bytes_store`, we reuse the existing entry |
| 47 | rather than storing duplicate bytes. This is common because submodules |
| 48 | often compile to identical artifacts (e.g., identical transformer layers |
| 49 | split on attn) |
| 50 | """ |
| 51 | |
| 52 | def __init__(self) -> None: |
| 53 | # dict from submodule name to byte hash |
| 54 | self.submodule_bytes: dict[str, str] = {} |
| 55 | # dict from byte hash to bytes |
| 56 | self.submodule_bytes_store: dict[str, bytes] = {} |
| 57 | # dict from byte hash to loaded module |
| 58 | self.loaded_submodule_store: dict[str, Any] = {} |
| 59 | |
| 60 | def insert(self, submod_name: str, shape: str, entry: bytes) -> None: |
| 61 | hasher = hashlib.sha256() |
| 62 | hasher.update(entry) |
| 63 | hex_digest = hasher.hexdigest() |
| 64 | self.submodule_bytes[f"{submod_name}_{shape}"] = hex_digest |
| 65 | if hex_digest not in self.submodule_bytes_store: |
| 66 | self.submodule_bytes_store[hex_digest] = entry |
| 67 | compilation_counter.num_compiled_artifacts_saved += 1 |
| 68 | logger.debug( |
| 69 | "inserting new artifact for submod %s with shape %s " |
| 70 | "(%s bytes) at hash %s", |
| 71 | submod_name, |
| 72 | shape, |
| 73 | len(entry), |
| 74 | hex_digest, |
| 75 | ) |
| 76 | else: |
| 77 | logger.debug( |
| 78 | "reusing existing cache artifact for submod %s " |
| 79 | "with shape %s (%s bytes) at hash %s", |
| 80 | submod_name, |
| 81 | shape, |
| 82 | len(entry), |
| 83 | hex_digest, |
| 84 | ) |
| 85 | |
| 86 | def get(self, submod_name: str, shape: str) -> bytes: |
| 87 | logger.debug( |
| 88 | "getting artifact for submod %s with shape %s", |
| 89 | submod_name, |
| 90 | shape, |
| 91 | ) |
| 92 | return self.submodule_bytes_store[ |
| 93 | self.submodule_bytes[f"{submod_name}_{shape}"] |
| 94 | ] |
| 95 | |
| 96 | def get_loaded(self, submod_name: str, shape: str) -> Any: |
| 97 | logger.debug( |
| 98 | "getting artifact for submod %s with shape %s", |
| 99 | submod_name, |
| 100 | shape, |
| 101 | ) |
| 102 | return self.loaded_submodule_store[ |
| 103 | self.submodule_bytes[f"{submod_name}_{shape}"] |
| 104 | ] |
| 105 | |
| 106 | def size_bytes(self) -> int: |
| 107 | return sum(len(entry) for entry in self.submodule_bytes_store.values()) |
| 108 | |
| 109 | def num_artifacts(self) -> int: |
| 110 | return len(self.submodule_bytes_store) |
| 111 | |
| 112 | def num_entries(self) -> int: |
| 113 | return len(self.submodule_bytes) |
| 114 | |
| 115 | def submodule_names(self) -> list[str]: |
| 116 | # get unique "{submod_name}" from "{submod_name}_{shape}", preserving order |
| 117 | names = [cache_key.rsplit("_", 1)[0] for cache_key in self.submodule_bytes] |
| 118 | return list(dict.fromkeys(names)) |
| 119 | |
| 120 | def load_all(self) -> None: |
| 121 | import concurrent.futures |
| 122 | |
| 123 | # check already loaded |
| 124 | if len(self.loaded_submodule_store) == len(self.submodule_bytes_store): |
| 125 | return |
| 126 | |
| 127 | from torch._inductor.standalone_compile import AOTCompiledArtifact |
| 128 | |
| 129 | def _load_entry(entry_bytes: bytes) -> AOTCompiledArtifact: |
| 130 | entry = pickle.loads(entry_bytes) |
| 131 | compilation_counter.num_compiled_artifacts_loaded += 1 |
| 132 | return AOTCompiledArtifact.deserialize(entry) |
| 133 | |
| 134 | entries = list(self.submodule_bytes_store.values()) |
| 135 | if current_platform.is_rocm(): |
| 136 | # Deserializing an artifact loads its compiled Triton kernels, which |
| 137 | # on ROCm ends up in hipModuleLoad. Loading modules concurrently |
| 138 | # deadlocks on the dynamic linker lock, since rocprofiler-sdk calls |
| 139 | # dl_iterate_phdr from inside the load, so load one at a time. |
| 140 | loaded_entries = [_load_entry(entry) for entry in entries] |
| 141 | else: |
| 142 | with concurrent.futures.ThreadPoolExecutor() as executor: |
| 143 | loaded_entries = list(executor.map(_load_entry, entries)) |
| 144 | |
| 145 | for i, k in enumerate(self.submodule_bytes_store.keys()): |
| 146 | self.loaded_submodule_store[k] = loaded_entries[i] |
| 147 | |
| 148 | logger.debug("loaded all %s submodules", self.num_artifacts()) |
| 149 | |
| 150 | def __getstate__(self) -> dict[str, dict[str, str] | dict[str, bytes]]: |
| 151 | return { |
| 152 | "submodule_bytes": self.submodule_bytes, |
| 153 | "submodule_bytes_store": self.submodule_bytes_store, |
| 154 | } |
| 155 | |
| 156 | def __setstate__(self, state: dict[str, dict[str, Any]]) -> None: |
| 157 | self.submodule_bytes = state["submodule_bytes"] |
| 158 | self.submodule_bytes_store = state["submodule_bytes_store"] |
| 159 | self.loaded_submodule_store = {} |
| 160 | |
| 161 | |
| 162 | @contextlib.contextmanager |
| 163 | def patch_pytree_map_over_slice(): |
| 164 | pytree._private_register_pytree_node( |
| 165 | slice, lambda x: ([x.start, x.stop, x.step], None), lambda x, c: slice(*x) |
| 166 | ) |
| 167 | |
| 168 | try: |
| 169 | yield |
| 170 | finally: |
| 171 | pytree._deregister_pytree_node(slice) |
| 172 | |
| 173 | |
| 174 | class VllmSerializableFunction(SerializableCallable): # type: ignore[misc] |
| 175 | """A wrapper around a compiled function by vllm. It will forward the tensor |
| 176 | inputs to the compiled function and return the result. |
| 177 | It also implements a serialization interface to support PyTorch's precompile |
| 178 | with custom backend, so that we can save and load the compiled function on |
| 179 | disk. There's no need to wrap around the compiled function if we don't want |
| 180 | to serialize them in particular cases. |
| 181 | Right now serialization for the custom backend is done via |
| 182 | serializing the Dynamo fx graph plus example inputs. |
| 183 | """ |
| 184 | |
| 185 | def __init__( |
| 186 | self, |
| 187 | graph_module: torch.fx.GraphModule | bytes, |
| 188 | example_inputs: Sequence[Any], |
| 189 | prefix: str, |
| 190 | optimized_call: Callable[..., Any], |
| 191 | is_encoder: bool = False, |
| 192 | vllm_backend: Any | None = None, |
| 193 | sym_tensor_indices: list[int] | None = None, |
| 194 | aot_autograd_config: dict[str, Any] | None = None, |
| 195 | execution_code: str | None = None, |
| 196 | submod_names: list[str] | None = None, |
| 197 | consts: list[Any] | None = None, |
| 198 | ) -> None: |
| 199 | self.graph_module = graph_module |
| 200 | self.example_inputs = example_inputs |
| 201 | self.prefix = prefix |
| 202 | self.optimized_call = optimized_call |
| 203 | self.is_encoder = is_encoder |
| 204 | self.shape_env = None |
| 205 | self.vllm_backend = vllm_backend |
| 206 | self.sym_tensor_indices = sym_tensor_indices |
| 207 | self.execution_code = execution_code |
| 208 | self.submod_names = submod_names |
| 209 | self.consts = consts |
| 210 | self._fake_mode: Any | None = None |
| 211 | |
| 212 | import torch._functorch.config as functorch_config |
| 213 | |
| 214 | self.aot_autograd_config = ( |
| 215 | aot_autograd_config or functorch_config.save_config_portable() |
| 216 | ) |
| 217 | sym_input = next( |
| 218 | (i for i in self.example_inputs if isinstance(i, torch.SymInt)), None |
| 219 | ) |
| 220 | if sym_input is not None: |
| 221 | self.shape_env = sym_input.node.shape_env |
| 222 | |
| 223 | def __call__(self, *args: Any, **kwargs: Any) -> Any: |
| 224 | return self.optimized_call(*args, **kwargs) |
| 225 | |
| 226 | @classmethod |
| 227 | def serialize_graph_module(cls, graph_module: torch.fx.GraphModule) -> bytes: |
| 228 | import sympy |
| 229 | |
| 230 | graph_reducer_override = GraphPickler.reducer_override |
| 231 | |
| 232 | def _graph_reducer_override( |
| 233 | self: GraphPickler, obj: Any |
| 234 | ) -> tuple[Callable[..., Any], tuple[Any, ...]] | Any: |
| 235 | if ( |
| 236 | inspect.isclass(obj) |
| 237 | and issubclass(obj, sympy.Function) |
| 238 | and hasattr(obj, "_torch_unpickler") |
| 239 | ): |
| 240 | return obj._torch_unpickler, (obj._torch_handler_name,) |
| 241 | if isinstance(obj, FakeTensorMode): |
| 242 | return type(None), () |
| 243 | return graph_reducer_override(self, obj) |
| 244 | |
| 245 | with ( |
| 246 | patch.object(GraphPickler, "reducer_override", _graph_reducer_override), |
| 247 | patch_pytree_map_over_slice(), |
| 248 | ): |
| 249 | return GraphPickler.dumps(graph_module, Options(ops_filter=None)) |
| 250 | |
| 251 | @classmethod |
| 252 | def deserialize_graph_module( |
| 253 | cls, data: bytes, fake_mode: FakeTensorMode |
| 254 | ) -> torch.fx.GraphModule: |
| 255 | with patch_pytree_map_over_slice(): |
| 256 | return GraphPickler.loads(data, fake_mode) |
| 257 | |
| 258 | @classmethod |
| 259 | def serialize_compile_artifacts( |
| 260 | cls, compiled_fn: "VllmSerializableFunction" |
| 261 | ) -> bytes: |
| 262 | state = compiled_fn.__dict__.copy() |
| 263 | state.pop("optimized_call") |
| 264 | state.pop("shape_env") |
| 265 | state.pop("vllm_backend", None) |
| 266 | state.pop("_fake_mode", None) |
| 267 | for node in state["graph_module"].graph.nodes: |
| 268 | node.meta.pop("source_fn_stack", None) |
| 269 | node.meta.pop("nn_module_stack", None) |
| 270 | for name, submod in state["graph_module"].named_children(): |
| 271 | if hasattr(submod, "graph"): |
| 272 | for node in submod.graph.nodes: |
| 273 | node.meta.pop("source_fn_stack", None) |
| 274 | node.meta.pop("nn_module_stack", None) |
| 275 | |
| 276 | if state.get("sym_tensor_indices"): |
| 277 | # put tensor inputs on meta device since their data |
| 278 | # isn't needed, yet we need the meta for make_copy_and_call |
| 279 | state["example_inputs"] = pytree.tree_map_only( |
| 280 | torch.Tensor, |
| 281 | lambda inp: torch.empty_like(inp, device="meta"), |
| 282 | state["example_inputs"], |
| 283 | ) |
| 284 | else: |
| 285 | # mask off all tensor inputs since they are large and not needed. |
| 286 | state["example_inputs"] = pytree.tree_map_only( |
| 287 | torch.Tensor, |
| 288 | lambda inp: torch.empty_like(inp, device="meta"), |
| 289 | state["example_inputs"], |
| 290 | ) |
| 291 | |
| 292 | state["graph_module"] = cls.serialize_graph_module(state["graph_module"]) |
| 293 | state["example_inputs"] = GraphPickler.dumps(state["example_inputs"]) |
| 294 | |
| 295 | if compiled_fn.vllm_backend: |
| 296 | ( |
| 297 | standalone_compile_artifacts, |
| 298 | sym_shape_indices_map, |
| 299 | returns_tuple_map, |
| 300 | ) = compiled_fn.vllm_backend.collect_standalone_compile_artifacts() |
| 301 | state["standalone_compile_artifacts"] = standalone_compile_artifacts |
| 302 | state["sym_shape_indices_map"] = sym_shape_indices_map |
| 303 | state["returns_tuple_map"] = returns_tuple_map |
| 304 | return pickle.dumps(state) |
| 305 | |
| 306 | @classmethod |
| 307 | def deserialize_compile_artifacts(cls, data: bytes) -> "VllmSerializableFunction": |
| 308 | from torch._guards import TracingContext, tracing |
| 309 | from torch.fx.experimental.symbolic_shapes import ShapeEnv |
| 310 | |
| 311 | state = pickle.loads(data) |
| 312 | fake_mode = FakeTensorMode(shape_env=ShapeEnv()) |
| 313 | |
| 314 | state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode) |
| 315 | |
| 316 | standalone_compile_artifacts = state.pop("standalone_compile_artifacts", None) |
| 317 | sym_shape_indices_map = state.pop("sym_shape_indices_map", {}) |
| 318 | returns_tuple_map = state.pop("returns_tuple_map", {}) |
| 319 | |
| 320 | saved_aot_autograd_config = state["aot_autograd_config"] |
| 321 | if saved_aot_autograd_config is not None: |
| 322 | functorch_ctx = torch._functorch.config.patch(saved_aot_autograd_config) |
| 323 | else: |
| 324 | functorch_ctx = contextlib.nullcontext() |
| 325 | |
| 326 | if envs.VLLM_USE_MEGA_AOT_ARTIFACT: |
| 327 | assert standalone_compile_artifacts is not None |
| 328 | submod_names = standalone_compile_artifacts.submodule_names() |
| 329 | num_submods = len(submod_names) |
| 330 | num_artifacts = standalone_compile_artifacts.num_artifacts() |
| 331 | |
| 332 | with functorch_ctx: |
| 333 | fn = reconstruct_serializable_fn_from_mega_artifact( |
| 334 | state=state, |
| 335 | standalone_compile_artifacts=standalone_compile_artifacts, |
| 336 | vllm_config=get_current_vllm_config(), |
| 337 | sym_shape_indices_map=sym_shape_indices_map, |
| 338 | returns_tuple_map=returns_tuple_map, |
| 339 | fake_mode=fake_mode, |
| 340 | ) |
| 341 | |
| 342 | logger.info( |
| 343 | "reconstructed serializable fn from standalone compile " |
| 344 | "artifacts. num_artifacts=%d num_submods=%d", |
| 345 | num_artifacts, |
| 346 | num_submods, |
| 347 | ) |
| 348 | |
| 349 | return fn |
| 350 | |
| 351 | state["graph_module"] = cls.deserialize_graph_module( |
| 352 | state["graph_module"], fake_mode |
| 353 | ) |
| 354 | state["graph_module"].recompile() |
| 355 | |
| 356 | # Fall back to standard VllmBackend. |
| 357 | # Use a lazy closure: the backend needs traced_files for cache |
| 358 | # dir computation, but those are only populated after |
| 359 | # _verify_source_unchanged runs in decorators.py (which happens |
| 360 | # after deserialization completes). |
| 361 | from vllm.compilation.backends import VllmBackend |
| 362 | |
| 363 | is_encoder = state.get("is_encoder", False) |
| 364 | vllm_config = get_current_vllm_config() |
| 365 | compile_inputs = list(state["example_inputs"]) |
| 366 | |
| 367 | def optimized_call(*example_inputs: Any) -> Any: |
| 368 | vllm_backend: VllmBackend = VllmBackend( |
| 369 | vllm_config, state["prefix"], is_encoder |
| 370 | ) |
| 371 | with tracing(TracingContext(fake_mode)), functorch_ctx: |
| 372 | fn.optimized_call = vllm_backend( |
| 373 | state["graph_module"], compile_inputs |
| 374 | ).optimized_call |
| 375 | fn.vllm_backend = vllm_backend |
| 376 | return fn.optimized_call(*example_inputs) |
| 377 | |
| 378 | fn = cls(**state, optimized_call=optimized_call) |
| 379 | fn._fake_mode = fake_mode |
| 380 | return fn |
| 381 | |
| 382 | def finalize_loading(self, vllm_config: VllmConfig) -> None: |
| 383 | """Eagerly initialize the compiled backend and perform all loading. |
| 384 | |
| 385 | Must be called after _verify_source_unchanged has populated |
| 386 | compilation_config.traced_files, which is needed for cache dir |
| 387 | computation. |
| 388 | """ |
| 389 | if self._fake_mode is None: |
| 390 | return # Already finalized, or mega path (no _fake_mode set) |
| 391 | |
| 392 | from torch._guards import TracingContext, tracing |
| 393 | |
| 394 | from vllm.compilation.backends import VllmBackend |
| 395 | |
| 396 | saved_aot_autograd_config = self.aot_autograd_config |
| 397 | if saved_aot_autograd_config is not None: |
| 398 | functorch_ctx = torch._functorch.config.patch(saved_aot_autograd_config) |
| 399 | else: |
| 400 | functorch_ctx = contextlib.nullcontext() |
| 401 | |
| 402 | vllm_backend = VllmBackend(vllm_config, self.prefix, self.is_encoder) |
| 403 | with tracing(TracingContext(self._fake_mode)), functorch_ctx: |
| 404 | result = vllm_backend(self.graph_module, list(self.example_inputs)) |
| 405 | self.optimized_call = result.optimized_call |
| 406 | self.vllm_backend = vllm_backend |
| 407 | |
| 408 | self._fake_mode = None |
| 409 | |
| 410 | @property |
| 411 | def co_name(self) -> Literal["VllmSerializableFunction"]: |
| 412 | """Used for depyf debugging.""" |
| 413 | return "VllmSerializableFunction" |
| 414 | |
| 415 | |
| 416 | def reconstruct_serializable_fn_from_mega_artifact( |
| 417 | state: dict[str, Any], |
| 418 | standalone_compile_artifacts: "StandaloneCompiledArtifacts", |
| 419 | vllm_config: VllmConfig, |
| 420 | sym_shape_indices_map: dict[str, list[int]], |
| 421 | returns_tuple_map: dict[str, bool], |
| 422 | fake_mode: FakeTensorMode, |
| 423 | ) -> "VllmSerializableFunction": |
| 424 | """Construct a VllmSerializableFunction from cached inductor artifacts. |
| 425 | |
| 426 | This function reconstructs a callable model from pre-compiled inductor |
| 427 | artifacts without re-running the compilation. It: |
| 428 | 1. Loads all cached artifacts |
| 429 | 2. Builds compiled callables for each submodule/shape |
| 430 | 3. Creates PiecewiseBackend instances that dispatch to cached artifacts |
| 431 | 4. Wraps with cudagraph if needed |
| 432 | 5. Returns the final VllmSerializableFunction |
| 433 | |
| 434 | Note: This function shares similar logic with PiecewiseCompileInterpreter |
| 435 | in backends.py. Both create PiecewiseBackend instances and wrap them with |
| 436 | cudagraph. The key difference is: |
| 437 | - this function: PiecewiseBackend receives pre-compiled runnables |
| 438 | (compiled_runnables is set, graph is None) |
| 439 | - PiecewiseCompileInterpreter: PiecewiseBackend receives the FX graph |
| 440 | to compile (graph is set, compiled_runnables is None) |
| 441 | |
| 442 | If modifying the backend creation/wrapping logic, consider updating both. |
| 443 | |
| 444 | Args: |
| 445 | fake_mode: FakeTensorMode used to unpickle the cached graph. |
| 446 | state: Deserialized state dict containing graph_module, example_inputs, |
| 447 | prefix, sym_tensor_indices, is_encoder, etc. |
| 448 | standalone_compile_artifacts: The StandaloneCompiledArtifacts containing |
| 449 | pre-compiled artifacts for each submodule/shape combination. |
| 450 | vllm_config: The vLLM configuration. |
| 451 | sym_shape_indices_map: Mapping from submod_name to sym_shape_indices. |
| 452 | returns_tuple_map: Mapping from submod_name to returns_tuple. |
| 453 | |
| 454 | Returns: |
| 455 | A VllmSerializableFunction that can be called directly. |
| 456 | |
| 457 | """ |
| 458 | from vllm.compilation.backends import ( |
| 459 | VllmBackend, |
| 460 | make_copy_and_call, |
| 461 | wrap_with_cudagraph_if_needed, |
| 462 | ) |
| 463 | from vllm.compilation.piecewise_backend import PiecewiseBackend |
| 464 | |
| 465 | prefix = state["prefix"] |
| 466 | is_encoder = state.get("is_encoder", False) |
| 467 | compilation_config = vllm_config.compilation_config |
| 468 | |
| 469 | standalone_compile_artifacts.load_all() |
| 470 | |
| 471 | piecewise_submod_names = standalone_compile_artifacts.submodule_names() |
| 472 | compiled_callables: dict[str, dict[str, Callable[..., Any]]] = {} |
| 473 | |
| 474 | for cache_key in standalone_compile_artifacts.submodule_bytes: |
| 475 | submod_name, shape_str = cache_key.rsplit("_", 1) |
| 476 | compiled_callables.setdefault(submod_name, {})[shape_str] = ( |
| 477 | standalone_compile_artifacts.get_loaded(submod_name, shape_str) |
| 478 | ) |
| 479 | |
| 480 | vllm_backend = VllmBackend(vllm_config, prefix, is_encoder) |
| 481 | dummy_cache_dir = os.path.join(envs.VLLM_CACHE_ROOT, "dummy_cache") |
| 482 | os.makedirs(dummy_cache_dir, exist_ok=True) |
| 483 | vllm_backend.compiler_manager.initialize_cache( |
| 484 | cache_dir=dummy_cache_dir, |
| 485 | disable_cache=True, |
| 486 | prefix=prefix, |
| 487 | ) |
| 488 | |
| 489 | # spot check that cached submodules exist in the graph structure |
| 490 | # if an old cache is used, this will fail but that's fine because |
| 491 | # we will just try this error and re-generate the new cache. |
| 492 | graph_children = set(state["submod_names"]) |
| 493 | missing = set(piecewise_submod_names) - graph_children |
| 494 | assert not missing, ( |
| 495 | f"artifacts reference submodules not in graph: {missing}. " |
| 496 | f"graph has: {sorted(graph_children)}" |
| 497 | ) |
| 498 | |
| 499 | submod_callables = {} |
| 500 | for i, submod_name in enumerate(piecewise_submod_names): |
| 501 | assert submod_name in sym_shape_indices_map and submod_name in returns_tuple_map |
| 502 | |
| 503 | sym_shape_indices = sym_shape_indices_map[submod_name] |
| 504 | returns_tuple = returns_tuple_map[submod_name] |
| 505 | runnables = compiled_callables[submod_name] |
| 506 | |
| 507 | piecewise_backend = PiecewiseBackend( |
| 508 | graph=None, # not needed for cached artifacts |
| 509 | vllm_config=vllm_config, |
| 510 | piecewise_compile_index=i, |
| 511 | total_piecewise_compiles=len(piecewise_submod_names), |
| 512 | sym_shape_indices=sym_shape_indices, |
| 513 | vllm_backend=vllm_backend, |
| 514 | returns_tuple=returns_tuple, |
| 515 | compiled_runnables=runnables, |
| 516 | ) |
| 517 | |
| 518 | is_first = i == 0 |
| 519 | is_last = i == len(piecewise_submod_names) - 1 |
| 520 | wrapped_backend = wrap_with_cudagraph_if_needed( |
| 521 | piecewise_backend, |
| 522 | vllm_config, |
| 523 | compilation_config, |
| 524 | is_first, |
| 525 | is_last, |
| 526 | ) |
| 527 | |
| 528 | submod_callables[submod_name] = wrapped_backend |
| 529 | logger.debug( |
| 530 | "Replaced submodule %s with piecewise backend from cache", |
| 531 | submod_name, |
| 532 | ) |
| 533 | |
| 534 | # Use codegen'd execution code if available, fall back to split_gm |
| 535 | execution_code = state.get("execution_code") |
| 536 | submod_names = state.get("submod_names") |
| 537 | if execution_code is not None and submod_names is not None: |
| 538 | consts = state.get("consts") |
| 539 | runtime_callable = compile_execution_fn( |
| 540 | execution_code, submod_callables, submod_names, consts |
| 541 | ) |
| 542 | else: |
| 543 | logger.warning( |
| 544 | "No execution code found, falling back to graph module execution." |
| 545 | ) |
| 546 | runtime_callable = GraphPickler.loads( |
| 547 | state["graph_module"], fake_mode=fake_mode |
| 548 | ) |
| 549 | |
| 550 | if compilation_config.cudagraph_copy_inputs: |
| 551 | sym_tensor_indices = state["sym_tensor_indices"] |
| 552 | input_buffers = [ |
| 553 | torch.empty_like( |
| 554 | state["example_inputs"][idx], device=vllm_config.device_config.device |
| 555 | ) |
| 556 | for idx in sym_tensor_indices |
| 557 | ] |
| 558 | optimized_call = make_copy_and_call( |
| 559 | sym_tensor_indices, input_buffers, runtime_callable |
| 560 | ) |
| 561 | else: |
| 562 | optimized_call = runtime_callable |
| 563 | |
| 564 | fn = VllmSerializableFunction( |
| 565 | **state, |
| 566 | optimized_call=optimized_call, |
| 567 | vllm_backend=None, |
| 568 | ) |
| 569 | return fn |
| 570 | |
| 571 | |
| 572 | def aot_compile_hash_factors(vllm_config: VllmConfig) -> list[str]: |
| 573 | factors = [] |
| 574 | # 0. factors come from the env, for example, The values of |
| 575 | # VLLM_PP_LAYER_PARTITION will affect the computation graph. |
| 576 | env_hash = hash_factors(envs.compile_factors()) |
| 577 | factors.append(env_hash) |
| 578 | |
| 579 | # 1. factors come from the vllm_config (it mainly summarizes how the |
| 580 | # model is created) |
| 581 | config_hash = vllm_config.compute_hash() |
| 582 | factors.append(config_hash) |
| 583 | |
| 584 | # 2. inductor factors if applicable |
| 585 | if envs.VLLM_USE_MEGA_AOT_ARTIFACT: |
| 586 | factors.extend(get_inductor_factors()) |
| 587 | |
| 588 | return factors |
| 589 | |
| 590 | |
| 591 | def _compute_code_hash_with_content(file_contents: dict[str, str]) -> str: |
| 592 | items = list(sorted(file_contents.items(), key=lambda x: x[0])) |
| 593 | hash_content = [] |
| 594 | for filepath, content in items: |
| 595 | hash_content.append(filepath) |
| 596 | if filepath == "<string>": |
| 597 | # This means the function was dynamically generated, with |
| 598 | # e.g. exec(). We can't actually check these. |
| 599 | continue |
| 600 | hash_content.append(content) |
| 601 | result: str = safe_hash( |
| 602 | "\n".join(hash_content).encode(), usedforsecurity=False |
| 603 | ).hexdigest() |
| 604 | return result |
| 605 | |
| 606 | |
| 607 | def _compute_code_hash(files: set[str]) -> str: |
| 608 | logger.debug( |
| 609 | "Traced files (to be considered for compilation cache):\n%s", "\n".join(files) |
| 610 | ) |
| 611 | file_contents = {} |
| 612 | for filepath in files: |
| 613 | # Skip files that don't exist (e.g., <string>, <frozen modules>, etc.) |
| 614 | if not os.path.isfile(filepath): |
| 615 | file_contents[filepath] = "" |
| 616 | else: |
| 617 | with open(filepath) as f: |
| 618 | file_contents[filepath] = f.read() |
| 619 | return _compute_code_hash_with_content(file_contents) |