| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import os |
| 5 | import sys |
| 6 | from abc import abstractmethod |
| 7 | from collections.abc import Callable, Generator |
| 8 | from contextlib import contextmanager, nullcontext |
| 9 | from types import CodeType |
| 10 | from typing import Any, ParamSpec, TypeVar |
| 11 | |
| 12 | import torch |
| 13 | |
| 14 | import vllm.envs as envs |
| 15 | from vllm.compilation.counter import compilation_counter |
| 16 | from vllm.config import ( |
| 17 | CompilationMode, |
| 18 | CUDAGraphMode, |
| 19 | VllmConfig, |
| 20 | get_current_vllm_config, |
| 21 | ) |
| 22 | from vllm.config.compilation import DynamicShapesType |
| 23 | from vllm.logger import init_logger |
| 24 | from vllm.utils.nvtx_pytorch_hooks import layerwise_nvtx_marker_context |
| 25 | |
| 26 | logger = init_logger(__name__) |
| 27 | |
| 28 | R = TypeVar("R") |
| 29 | P = ParamSpec("P") |
| 30 | |
| 31 | |
| 32 | @contextmanager |
| 33 | def _compilation_context() -> Generator[None, None, None]: |
| 34 | """Context manager for compilation settings. |
| 35 | |
| 36 | This manager sets higher dynamo cache limits for compilation. |
| 37 | (Needed for qwen2_5_vl see test_qwen2_5_vl_evs_functionality). |
| 38 | Generally a recompilation can happen whenever we use a new |
| 39 | backend instance in torch.compile. |
| 40 | """ |
| 41 | original_cache_size = torch._dynamo.config.cache_size_limit |
| 42 | original_accumulated_cache = torch._dynamo.config.accumulated_cache_size_limit |
| 43 | |
| 44 | try: |
| 45 | torch._dynamo.config.cache_size_limit = 2048 |
| 46 | torch._dynamo.config.accumulated_cache_size_limit = 8192 |
| 47 | yield |
| 48 | finally: |
| 49 | torch._dynamo.config.cache_size_limit = original_cache_size |
| 50 | torch._dynamo.config.accumulated_cache_size_limit = original_accumulated_cache |
| 51 | |
| 52 | |
| 53 | def compile_model_with_stock_torch( |
| 54 | model: torch.nn.Module, vllm_config: VllmConfig |
| 55 | ) -> None: |
| 56 | from vllm.env_override import _apply_constrain_to_fx_strides_patch |
| 57 | |
| 58 | _apply_constrain_to_fx_strides_patch() |
| 59 | backend = vllm_config.compilation_config.init_backend(vllm_config) |
| 60 | compilation_counter.stock_torch_compile_count += 1 |
| 61 | model.compile(fullgraph=True, backend=backend) |
| 62 | |
| 63 | |
| 64 | class TorchCompileWithNoGuardsWrapper: |
| 65 | """A wrapper class for torch.compile, it ensures that all guards are dropped |
| 66 | when CompilationMode is not CompilationMode.STOCK_TORCH_COMPILE. |
| 67 | When guards are dropped, the first time __call__ is invoked, a single |
| 68 | compilation is triggered. Dynamo should never be traced again after that |
| 69 | since we drop all guards. |
| 70 | """ |
| 71 | |
| 72 | def _call_with_optional_nvtx_range( |
| 73 | self, callable_fn: Callable[P, R], *args: P.args, **kwargs: P.kwargs |
| 74 | ) -> Any: |
| 75 | if self.layerwise_nvtx_tracing_enabled: |
| 76 | args_list = list(args) |
| 77 | kwargs_dict = dict(kwargs) |
| 78 | with layerwise_nvtx_marker_context( |
| 79 | "Torch Compiled Module (input):{}".format(self.__class__.__name__), |
| 80 | self, |
| 81 | in_tensor=args_list, |
| 82 | kwargs=kwargs_dict, |
| 83 | ) as ctx: |
| 84 | ctx.result = callable_fn(*args, **kwargs) |
| 85 | return ctx.result |
| 86 | return callable_fn(*args, **kwargs) |
| 87 | |
| 88 | def __init__( |
| 89 | self, |
| 90 | compile_prefix: str = "", |
| 91 | is_encoder: bool = False, |
| 92 | ) -> None: |
| 93 | self.compiled = False |
| 94 | self._compile_prefix = compile_prefix |
| 95 | self._is_encoder = is_encoder |
| 96 | |
| 97 | vllm_config = get_current_vllm_config() |
| 98 | self.vllm_config = vllm_config |
| 99 | mode = vllm_config.compilation_config.mode |
| 100 | self.layerwise_nvtx_tracing_enabled = ( |
| 101 | vllm_config.observability_config.enable_layerwise_nvtx_tracing |
| 102 | ) |
| 103 | if mode is None: |
| 104 | raise RuntimeError("Compilation mode cannot be NO_COMPILATION") |
| 105 | |
| 106 | backend = vllm_config.compilation_config.init_backend( |
| 107 | vllm_config, prefix=compile_prefix, is_encoder=is_encoder |
| 108 | ) |
| 109 | options = {} |
| 110 | |
| 111 | if isinstance(backend, str) and backend == "inductor": |
| 112 | options = vllm_config.compilation_config.inductor_compile_config |
| 113 | |
| 114 | self.first_compile = True |
| 115 | self.evaluate_guards = ( |
| 116 | vllm_config.compilation_config.dynamic_shapes_config.evaluate_guards |
| 117 | ) |
| 118 | |
| 119 | ds_type = vllm_config.compilation_config.dynamic_shapes_config.type |
| 120 | |
| 121 | if mode != CompilationMode.STOCK_TORCH_COMPILE: |
| 122 | # Drop all the guards. |
| 123 | if self.evaluate_guards: |
| 124 | assert not envs.VLLM_USE_BYTECODE_HOOK, ( |
| 125 | "compilation_config.dynamic_shapes_config.evaluate_guards " |
| 126 | "requires VLLM_USE_BYTECODE_HOOK=0. " |
| 127 | ) |
| 128 | assert ds_type != DynamicShapesType.UNBACKED, ( |
| 129 | "UNBACKED dynamic shapes do not add guards" |
| 130 | ) |
| 131 | |
| 132 | options["guard_filter_fn"] = lambda x: [ |
| 133 | entry.guard_type == "SHAPE_ENV" for entry in x |
| 134 | ] |
| 135 | else: |
| 136 | if hasattr(torch.compiler, "skip_all_guards_unsafe"): |
| 137 | # Torch 2.10+ provides skip_all_guards_unsafe |
| 138 | options["guard_filter_fn"] = torch.compiler.skip_all_guards_unsafe |
| 139 | else: |
| 140 | # Equivalent fallback for older PyTorch: skip all guards |
| 141 | options["guard_filter_fn"] = lambda x: [False for _ in x] |
| 142 | |
| 143 | compiled_ptr: Any = self.forward |
| 144 | # Validate that unbacked dynamic shapes require VLLM_USE_BYTECODE_HOOK=False |
| 145 | |
| 146 | # Apply the constrain_to_fx_strides patch before first compilation. |
| 147 | # This covers STOCK_TORCH_COMPILE and DYNAMO_ONCE paths. The VLLM |
| 148 | # compile paths call this from their own compile() methods too. |
| 149 | from vllm.env_override import _apply_constrain_to_fx_strides_patch |
| 150 | |
| 151 | _apply_constrain_to_fx_strides_patch() |
| 152 | |
| 153 | aot_context = nullcontext() |
| 154 | if envs.VLLM_USE_AOT_COMPILE: |
| 155 | if hasattr(torch._dynamo.config, "enable_aot_compile"): |
| 156 | aot_context = torch._dynamo.config.patch(enable_aot_compile=True) |
| 157 | else: |
| 158 | msg = "torch._dynamo.config.enable_aot_compile is not " |
| 159 | msg += "available. AOT compile is disabled and please " |
| 160 | msg += "upgrade PyTorch version to use AOT compile." |
| 161 | logger.warning(msg) |
| 162 | |
| 163 | with aot_context: |
| 164 | self._compiled_callable = torch.compile( |
| 165 | compiled_ptr, |
| 166 | fullgraph=True, |
| 167 | dynamic=False, |
| 168 | backend=backend, |
| 169 | options=options, |
| 170 | ) |
| 171 | |
| 172 | if envs.VLLM_USE_BYTECODE_HOOK and mode != CompilationMode.STOCK_TORCH_COMPILE: |
| 173 | self._bytecode_hook_handle = ( |
| 174 | torch._dynamo.convert_frame.register_bytecode_hook(self.bytecode_hook) |
| 175 | ) |
| 176 | self._compiled_bytecode: CodeType | None = None |
| 177 | |
| 178 | def aot_compile(self, *args: Any, **kwargs: Any) -> Any: |
| 179 | if not hasattr(self._compiled_callable, "aot_compile"): |
| 180 | raise RuntimeError( |
| 181 | "aot_compile is not supported by the current configuration. " |
| 182 | "Please make sure torch.compile is enabled with the latest " |
| 183 | f"version of PyTorch (current using torch: {torch.__version__})" |
| 184 | ) |
| 185 | return self._compiled_callable.aot_compile((args, kwargs)) |
| 186 | |
| 187 | def __call__(self, *args: Any, **kwargs: Any) -> Any: |
| 188 | if envs.VLLM_USE_BYTECODE_HOOK: |
| 189 | if ( |
| 190 | self.vllm_config.compilation_config.mode |
| 191 | == CompilationMode.STOCK_TORCH_COMPILE |
| 192 | ): |
| 193 | return self._compiled_callable(*args, **kwargs) |
| 194 | |
| 195 | if not self._compiled_bytecode: |
| 196 | # Make sure a compilation is triggered by clearing dynamo |
| 197 | # cache. |
| 198 | torch._dynamo.eval_frame.remove_from_cache(self.original_code_object()) |
| 199 | return self._call_with_optional_nvtx_range( |
| 200 | self._compiled_callable, *args, **kwargs |
| 201 | ) |
| 202 | else: |
| 203 | with self._dispatch_to_compiled_code(): |
| 204 | return self._call_with_optional_nvtx_range( |
| 205 | self.forward, *args, **kwargs |
| 206 | ) |
| 207 | else: |
| 208 | ctx = ( |
| 209 | nullcontext() |
| 210 | if self.first_compile or not self.evaluate_guards |
| 211 | else torch.compiler.set_stance("fail_on_recompile") |
| 212 | ) |
| 213 | self.first_compile = False |
| 214 | with _compilation_context(), ctx: |
| 215 | return self._call_with_optional_nvtx_range( |
| 216 | self._compiled_callable, *args, **kwargs |
| 217 | ) |
| 218 | |
| 219 | @abstractmethod |
| 220 | def forward(self, *args: Any, **kwargs: Any) -> Any: ... |
| 221 | |
| 222 | def original_code_object(self) -> CodeType: |
| 223 | """Return the original code object of the forward method.""" |
| 224 | return self.__class__.forward.__code__ |
| 225 | |
| 226 | def bytecode_hook(self, old_code: CodeType, new_code: CodeType) -> None: |
| 227 | """Hook to save the compiled bytecode for direct execution.""" |
| 228 | if old_code is not self.original_code_object(): |
| 229 | return |
| 230 | # code borrowed from https://github.com/thuml/depyf/blob/f4ad79fadee27ea113b4c75202db1eb1a11c0dbc/depyf/explain/enable_debugging.py#L25 |
| 231 | frame = sys._getframe() |
| 232 | while frame and frame.f_back: |
| 233 | frame = frame.f_back |
| 234 | code_name = frame.f_code.co_name |
| 235 | file_name = frame.f_code.co_filename.split(os.path.sep)[-1] |
| 236 | if code_name == "_compile" and file_name == "convert_frame.py": |
| 237 | break |
| 238 | frame = frame.f_locals["frame"] |
| 239 | assert frame.f_code == old_code |
| 240 | |
| 241 | if frame.f_locals["self"] is not self: |
| 242 | return |
| 243 | |
| 244 | self._compiled_bytecode = new_code |
| 245 | |
| 246 | path = self.vllm_config.compile_debug_dump_path() |
| 247 | if path: |
| 248 | decompiled_file = path / "transformed_code.py" |
| 249 | if not decompiled_file.exists(): |
| 250 | try: |
| 251 | # usually the decompilation will succeed for most models, |
| 252 | # as we guarantee a full-graph compilation in Dynamo. |
| 253 | # but there's no 100% guarantee, since decompliation is |
| 254 | # not a reversible process. |
| 255 | import depyf |
| 256 | |
| 257 | src = depyf.decompile(new_code) |
| 258 | |
| 259 | with open(decompiled_file, "w") as f: |
| 260 | f.write(src) |
| 261 | |
| 262 | logger.debug("Dynamo transformed code saved to %s", decompiled_file) |
| 263 | except Exception: |
| 264 | pass |
| 265 | |
| 266 | if ( |
| 267 | self.vllm_config.compilation_config.cudagraph_mode != CUDAGraphMode.NONE |
| 268 | and "update" in new_code.co_names |
| 269 | ): |
| 270 | import depyf |
| 271 | |
| 272 | src = depyf.decompile(new_code) |
| 273 | msg = ( |
| 274 | "Assigning / modifying buffers of nn.Module during forward pass is not " |
| 275 | "allowed when using cudagraph inside the compiler because it will " |
| 276 | "cause silent errors. Please use eager mode or fix the code. The " |
| 277 | "following code contains clues about which buffer is being modified " |
| 278 | f"(please search for the usage of the function `update`):\n{src}" |
| 279 | ) |
| 280 | raise RuntimeError(msg) |
| 281 | |
| 282 | def cleanup(self) -> None: |
| 283 | """Remove the bytecode hook registered by this instance.""" |
| 284 | handle = getattr(self, "_bytecode_hook_handle", None) |
| 285 | if handle is not None: |
| 286 | handle.remove() |
| 287 | |
| 288 | @contextmanager |
| 289 | def _dispatch_to_compiled_code(self) -> Generator[None, None, None]: |
| 290 | # noqa: E501 |
| 291 | """Context manager to dispatch to internally compiled code for torch<2.8. |
| 292 | Why does this work? Because Dynamo guarantees that the compiled |
| 293 | bytecode has exactly the same arguments, cell variables, and free |
| 294 | variables as the original code. Therefore we can directly switch |
| 295 | the code object in the function and call it. |
| 296 | |
| 297 | See https://dev-discuss.pytorch.org/t/what-is-the-relationship-requirement-among-original-bytecode-transformed-bytecode-and-bytecode-returned-by-hooks-in-dynamo/1693/7 for more details. |
| 298 | """ # noqa: E501 line too long |
| 299 | original = self.original_code_object() |
| 300 | assert self._compiled_bytecode is not None |
| 301 | self.__class__.forward.__code__ = self._compiled_bytecode |
| 302 | try: |
| 303 | yield |
| 304 | finally: |
| 305 | self.__class__.forward.__code__ = original |
| 306 | |
| 307 | |
| 308 | def reset_compile_wrapper(model: torch.nn.Module) -> None: |
| 309 | """Clean up compiled model and captured CUDA graphs for elastic EP.""" |
| 310 | if not isinstance(model, TorchCompileWithNoGuardsWrapper) and hasattr( |
| 311 | model, "model" |
| 312 | ): |
| 313 | model = model.model |
| 314 | if not isinstance(model, TorchCompileWithNoGuardsWrapper): |
| 315 | return |
| 316 | # model.do_not_compile is set by the @support_torch_compile decorator |
| 317 | if hasattr(model, "do_not_compile") and model.do_not_compile: |
| 318 | return |
| 319 | from vllm.compilation.counter import compilation_counter |
| 320 | |
| 321 | # reset the compilation counter |
| 322 | compilation_counter.num_models_seen = 0 |
| 323 | compilation_counter.num_graphs_seen = 0 |
| 324 | compilation_counter.num_piecewise_graphs_seen = 0 |
| 325 | compilation_counter.num_piecewise_capturable_graphs_seen = 0 |
| 326 | compilation_counter.num_backend_compilations = 0 |
| 327 | compilation_counter.num_gpu_runner_capture_triggers = 0 |
| 328 | compilation_counter.num_cudagraph_captured = 0 |
| 329 | compilation_counter.num_inductor_compiles = 0 |
| 330 | compilation_counter.num_eager_compiles = 0 |
| 331 | compilation_counter.num_cache_entries_updated = 0 |
| 332 | compilation_counter.num_compiled_artifacts_saved = 0 |
| 333 | compilation_counter.stock_torch_compile_count = 0 |
| 334 | compilation_counter.num_aot_compiles = 0 |
| 335 | compilation_counter.num_aot_artifacts_saved = 0 |
| 336 | compilation_counter.num_aot_artifacts_loaded = 0 |
| 337 | |
| 338 | # Clear the AOT compiled function so the model is forced to |
| 339 | # recompile on the next call. Without this, decorators.py |
| 340 | # __call__ uses the stale aot_compiled_fn whose torchinductor |
| 341 | # kernels have old parameters (expert_map size for example) |
| 342 | # baked in as compile-time constants. |
| 343 | if hasattr(model, "aot_compiled_fn"): |
| 344 | model.aot_compiled_fn = None |
| 345 | if hasattr(model, "was_aot_compile_fn_loaded_from_disk"): |
| 346 | model.was_aot_compile_fn_loaded_from_disk = False |
| 347 | |
| 348 | # Reset the cache_dir so VllmBackend recomputes the hash |
| 349 | # (data_parallel_size changed, so the config hash differs). |
| 350 | compilation_config = model.vllm_config.compilation_config |
| 351 | compilation_config.cache_dir = "" |
| 352 | compilation_config.local_cache_dir = "" |
| 353 | |
| 354 | model.__class__.forward.__code__ = model.original_code_object() |
| 355 | TorchCompileWithNoGuardsWrapper.__init__( |
| 356 | model, |
| 357 | compile_prefix=model._compile_prefix, |
| 358 | is_encoder=model._is_encoder, |
| 359 | ) |