| 1 | # SPDX-License-Identifier: Apache-2.0 |
| 2 | # SPDX-FileCopyrightText: Copyright contributors to the vLLM project |
| 3 | |
| 4 | import dataclasses |
| 5 | import io |
| 6 | import json |
| 7 | import pickle |
| 8 | from collections.abc import Callable |
| 9 | from pickle import Pickler |
| 10 | from typing import Any |
| 11 | |
| 12 | import torch._functorch.config |
| 13 | import torch.fx as fx |
| 14 | from torch._dynamo.utils import dynamo_timed |
| 15 | from torch._inductor.runtime.triton_heuristics import CachingAutotuner |
| 16 | from torch._logging._internal import trace_structured |
| 17 | |
| 18 | from vllm.compilation.backends import VllmBackend |
| 19 | from vllm.config import VllmConfig |
| 20 | from vllm.config.utils import Range |
| 21 | from vllm.logger import init_logger |
| 22 | |
| 23 | logger = init_logger(__name__) |
| 24 | |
| 25 | |
| 26 | def get_fake_args_from_graph(graph: fx.GraphModule) -> list[Any]: |
| 27 | """Get fake args directly from graph placeholder nodes.""" |
| 28 | fake_args = [] |
| 29 | for node in graph.graph.nodes: |
| 30 | if node.op == "placeholder": |
| 31 | fake_args.append(node.meta["example_value"]) |
| 32 | else: |
| 33 | break |
| 34 | return fake_args |
| 35 | |
| 36 | |
| 37 | def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]: |
| 38 | """Create Fake example inputs with symbolic dims replaced by a concrete size. |
| 39 | |
| 40 | Used for single-size compilation where we need concrete-shaped inputs. |
| 41 | The Dynamo-captured graph gives us example inputs with SymInts in them. |
| 42 | """ |
| 43 | from torch._prims_common import compute_required_storage_length |
| 44 | from torch._subclasses.fake_tensor import FakeTensorMode |
| 45 | from torch.fx.experimental.symbolic_shapes import ShapeEnv, is_symbolic |
| 46 | |
| 47 | def concretize(sym_val: Any) -> int: |
| 48 | """Replace all symbolic variables in a SymInt expression with size.""" |
| 49 | if not is_symbolic(sym_val): |
| 50 | return int(sym_val) |
| 51 | expr = sym_val.node.expr |
| 52 | return int(expr.subs({s: size for s in expr.free_symbols})) |
| 53 | |
| 54 | fake_mode = FakeTensorMode(shape_env=ShapeEnv()) |
| 55 | |
| 56 | args: list[Any] = [] |
| 57 | with fake_mode: |
| 58 | for node in graph.graph.nodes: |
| 59 | if node.op != "placeholder": |
| 60 | break |
| 61 | val = node.meta["example_value"] |
| 62 | if isinstance(val, torch.SymInt): |
| 63 | args.append(concretize(val)) |
| 64 | elif isinstance(val, torch.Tensor): |
| 65 | new_shape = tuple(concretize(d) for d in val.shape) |
| 66 | new_strides = tuple(concretize(s) for s in val.stride()) |
| 67 | new_storage_offset = concretize(val.storage_offset()) |
| 68 | needed_size = compute_required_storage_length( |
| 69 | new_shape, new_strides, new_storage_offset |
| 70 | ) |
| 71 | t = torch.empty(needed_size, dtype=val.dtype, device=val.device) |
| 72 | t = t.as_strided(new_shape, new_strides, new_storage_offset) |
| 73 | args.append(t) |
| 74 | else: |
| 75 | args.append(val) |
| 76 | return args |
| 77 | |
| 78 | |
| 79 | @dataclasses.dataclass |
| 80 | class RangeEntry: |
| 81 | compile_range: Range |
| 82 | compiled: bool = False |
| 83 | runnable: Callable[..., Any] = None # type: ignore |
| 84 | |
| 85 | |
| 86 | class PiecewiseBackend: |
| 87 | def __init__( |
| 88 | self, |
| 89 | graph: fx.GraphModule | None, |
| 90 | vllm_config: VllmConfig, |
| 91 | piecewise_compile_index: int, |
| 92 | total_piecewise_compiles: int, |
| 93 | sym_shape_indices: list[int], |
| 94 | vllm_backend: VllmBackend, |
| 95 | returns_tuple: bool, |
| 96 | compiled_runnables: dict[str, Callable[..., Any]] | None = None, |
| 97 | submod_name: str = "", |
| 98 | ): |
| 99 | """The backend for piecewise compilation. |
| 100 | It mainly handles the compilation of static shapes and |
| 101 | dispatching based on runtime shape. |
| 102 | |
| 103 | We will compile `self.graph` once for the general shape, |
| 104 | and then compile for different shapes specified in |
| 105 | `compilation_config.compile_sizes`. |
| 106 | |
| 107 | This class supports two mutually exclusive modes: |
| 108 | 1. Compilation (graph is set, compiled_runnables is None): |
| 109 | Used during initial compilation when we have the FX graph |
| 110 | and need to compile it for each shape range. |
| 111 | 2. Precompilation (graph is None, compiled_runnables is set): |
| 112 | Used when loading from cache/AOT artifacts where we already |
| 113 | have pre-compiled callables and don't need the original graph. |
| 114 | |
| 115 | Exactly one of graph or compiled_runnables must be provided. |
| 116 | """ |
| 117 | assert bool(graph is not None) ^ bool(compiled_runnables is not None), ( |
| 118 | "exactly one of graph and compiled_runnables should be set." |
| 119 | ) |
| 120 | |
| 121 | self.graph = graph |
| 122 | self.vllm_config = vllm_config |
| 123 | self.compilation_config = vllm_config.compilation_config |
| 124 | self.piecewise_compile_index = piecewise_compile_index |
| 125 | self.total_piecewise_compiles = total_piecewise_compiles |
| 126 | self.vllm_backend = vllm_backend |
| 127 | self.compiled_runnables = compiled_runnables |
| 128 | self.submod_name = submod_name |
| 129 | |
| 130 | self.is_first_graph = piecewise_compile_index == 0 |
| 131 | self.is_last_graph = piecewise_compile_index == total_piecewise_compiles - 1 |
| 132 | |
| 133 | self.is_full_graph = total_piecewise_compiles == 1 |
| 134 | self.is_encoder_compilation = vllm_backend.is_encoder |
| 135 | |
| 136 | self.compile_ranges = self.compilation_config.get_compile_ranges() |
| 137 | if self.is_encoder_compilation: |
| 138 | # For encoder compilation we use the max int32 value |
| 139 | # to set the upper bound of the compile ranges |
| 140 | max_int32 = 2**31 - 1 |
| 141 | last_compile_range = self.compile_ranges[-1] |
| 142 | assert ( |
| 143 | last_compile_range.end |
| 144 | == vllm_config.scheduler_config.max_num_batched_tokens |
| 145 | ) |
| 146 | self.compile_ranges[-1] = Range( |
| 147 | start=last_compile_range.start, end=max_int32 |
| 148 | ) |
| 149 | |
| 150 | log_string = f"PiecewiseBackend: compile_ranges: {self.compile_ranges}" |
| 151 | logger.debug_once(log_string) |
| 152 | |
| 153 | self.compile_sizes = self.compilation_config.compile_sizes |
| 154 | log_string = f"PiecewiseBackend: compile_sizes: {self.compile_sizes}" |
| 155 | logger.debug_once(log_string) |
| 156 | |
| 157 | self.sym_shape_indices = sym_shape_indices |
| 158 | self.returns_tuple = returns_tuple |
| 159 | |
| 160 | # the entries for ranges that we need to either |
| 161 | self.range_entries: dict[Range, RangeEntry] = {} |
| 162 | |
| 163 | # We only keep compilation management inside this class directly. |
| 164 | if self.compile_sizes is not None: |
| 165 | for size in self.compile_sizes: |
| 166 | if isinstance(size, str): |
| 167 | assert size == "cudagraph_capture_sizes" |
| 168 | raise NotImplementedError( |
| 169 | "cudagraph_capture_sizes not supported in compile_sizes." |
| 170 | "This should be handled in `post_init_cudagraph_sizes`." |
| 171 | ) |
| 172 | else: |
| 173 | assert isinstance(size, int) |
| 174 | range = Range(start=size, end=size) |
| 175 | if range not in self.compile_ranges: |
| 176 | self.range_entries[range] = RangeEntry( |
| 177 | compile_range=range, |
| 178 | ) |
| 179 | |
| 180 | for range in self.compile_ranges: |
| 181 | self.range_entries[range] = RangeEntry( |
| 182 | compile_range=range, |
| 183 | ) |
| 184 | |
| 185 | # Track whether we've logged the graph for this subgraph (only log once) |
| 186 | self._graph_logged = False |
| 187 | |
| 188 | if self.graph is not None: |
| 189 | self.compile_all_ranges() |
| 190 | else: |
| 191 | self.load_all_ranges() |
| 192 | |
| 193 | def get_compiled_graph_wrapper( |
| 194 | self, compiled_graph: Callable[..., Any] |
| 195 | ) -> Callable[..., Any]: |
| 196 | def compiled_graph_wrapper(*args: Any) -> Any: |
| 197 | graph_output = compiled_graph(*args) |
| 198 | # unpack the tuple if needed |
| 199 | # TODO(rzou): the implication is that we're not |
| 200 | # reading the python bytecode correctly in vLLM? |
| 201 | if self.returns_tuple or not isinstance(graph_output, (tuple, list)): |
| 202 | return graph_output |
| 203 | else: |
| 204 | return graph_output[0] |
| 205 | |
| 206 | return compiled_graph_wrapper |
| 207 | |
| 208 | def to_bytes(self) -> dict[str, bytes]: |
| 209 | class StandaloneCompiledArtifactsPickler(Pickler): |
| 210 | def reducer_override(self, obj: object) -> Any: |
| 211 | if isinstance(obj, CachingAutotuner): |
| 212 | obj.prepare_for_pickle() |
| 213 | return pickle.loads, ( |
| 214 | pickle.dumps( |
| 215 | obj, |
| 216 | ), |
| 217 | ) |
| 218 | return NotImplemented |
| 219 | |
| 220 | def serialize(fn: Callable[..., Any]) -> bytes: |
| 221 | assert hasattr(fn, "serialize"), "fn must have serialize method" |
| 222 | with torch._functorch.config.patch("bundled_autograd_cache", True): |
| 223 | entry = fn.serialize() |
| 224 | |
| 225 | f = io.BytesIO() |
| 226 | StandaloneCompiledArtifactsPickler(f).dump(entry) |
| 227 | result = f.getvalue() |
| 228 | return result |
| 229 | |
| 230 | out = {} |
| 231 | |
| 232 | for range_key, entry in self.range_entries.items(): |
| 233 | if not entry.compiled: |
| 234 | logger.debug( |
| 235 | "entry with range %s not compiled, so cannot get its bytes", |
| 236 | range_key, |
| 237 | ) |
| 238 | continue |
| 239 | if hasattr(entry.runnable, "serialize"): |
| 240 | out[str(range_key)] = serialize(entry.runnable) |
| 241 | |
| 242 | return out |
| 243 | |
| 244 | def compile_all_ranges(self) -> None: |
| 245 | """Compile all range entries for this piecewise subgraph up front.""" |
| 246 | assert self.graph is not None, ( |
| 247 | "Cannot compile without a graph. " |
| 248 | "When loading from cache/AOT artifacts, " |
| 249 | "compile_all_ranges should not be called." |
| 250 | ) |
| 251 | |
| 252 | for range_entry in self.range_entries.values(): |
| 253 | if range_entry.compiled: |
| 254 | continue |
| 255 | |
| 256 | self._log_compile_start(range_entry.compile_range) |
| 257 | |
| 258 | if range_entry.compile_range.is_single_size(): |
| 259 | args_list = create_concrete_args( |
| 260 | self.graph, range_entry.compile_range.start |
| 261 | ) |
| 262 | else: |
| 263 | args_list = get_fake_args_from_graph(self.graph) |
| 264 | |
| 265 | range_entry.runnable = self.vllm_backend.compiler_manager.compile( |
| 266 | self.graph, |
| 267 | args_list, |
| 268 | self.vllm_backend.inductor_config, |
| 269 | self.compilation_config, |
| 270 | compile_range=range_entry.compile_range, |
| 271 | graph_index=self.piecewise_compile_index, |
| 272 | num_graphs=self.total_piecewise_compiles, |
| 273 | is_encoder=self.vllm_backend.is_encoder, |
| 274 | ) |
| 275 | |
| 276 | range_entry.compiled = True |
| 277 | |
| 278 | @dynamo_timed("vllm_log_compile_start_torch_trace_only") |
| 279 | def _log_compile_start(self, compile_range: Range): |
| 280 | """Log compilation event for TORCH_TRACE/tlparse.""" |
| 281 | is_cudagraph_size = ( |
| 282 | self.compile_sizes is not None and compile_range.start in self.compile_sizes |
| 283 | ) |
| 284 | subgraph_index = self.piecewise_compile_index |
| 285 | submod_name = self.submod_name |
| 286 | trace_structured( |
| 287 | "artifact", |
| 288 | metadata_fn=lambda: { |
| 289 | "name": "vllm_piecewise_compile_start", |
| 290 | "encoding": "json", |
| 291 | }, |
| 292 | payload_fn=lambda: json.dumps( |
| 293 | { |
| 294 | "piecewise_index": subgraph_index, |
| 295 | "submod_name": submod_name, |
| 296 | "total_piecewise_compiles": self.total_piecewise_compiles, |
| 297 | "compile_range_start": compile_range.start, |
| 298 | "compile_range_end": compile_range.end, |
| 299 | "is_single_size": compile_range.is_single_size(), |
| 300 | "is_cudagraph_capture_size": is_cudagraph_size, |
| 301 | } |
| 302 | ), |
| 303 | ) |
| 304 | |
| 305 | # Log the subgraph graph dump only once per subgraph (not per size) |
| 306 | # to reduce log file size. The graph code is the same for all sizes. |
| 307 | if not self._graph_logged: |
| 308 | self._graph_logged = True |
| 309 | assert self.graph is not None |
| 310 | trace_structured( |
| 311 | "graph_dump", |
| 312 | metadata_fn=lambda: { |
| 313 | "name": f"vllm_{submod_name}", |
| 314 | }, |
| 315 | payload_fn=lambda: self.graph.print_readable(print_output=False), |
| 316 | ) |
| 317 | |
| 318 | def load_all_ranges(self) -> None: |
| 319 | """Load all pre-compiled runnables for this piecewise subgraph. |
| 320 | |
| 321 | Called during warm start to wrap all cached compiled_runnables |
| 322 | into range_entry.runnable up front, analogous to compile_all_ranges() |
| 323 | for the cold start path. |
| 324 | """ |
| 325 | assert self.compiled_runnables is not None, ( |
| 326 | "load_all_ranges should only be called when compiled_runnables " |
| 327 | "is set (warm start / cache loading path)." |
| 328 | ) |
| 329 | for range_entry in self.range_entries.values(): |
| 330 | if range_entry.compiled: |
| 331 | continue |
| 332 | key = str(range_entry.compile_range) |
| 333 | assert key in self.compiled_runnables, ( |
| 334 | f"Missing compiled runnable for range {range_entry.compile_range}. " |
| 335 | f"Available keys: {list(self.compiled_runnables.keys())}" |
| 336 | ) |
| 337 | range_entry.runnable = self.get_compiled_graph_wrapper( |
| 338 | self.compiled_runnables[key] |
| 339 | ) |
| 340 | range_entry.compiled = True |
| 341 | |
| 342 | def _find_range_for_shape(self, runtime_shape: int) -> RangeEntry | None: |
| 343 | # First we try to find the range entry for the concrete compile size |
| 344 | # If not found, we search for the range entry |
| 345 | # that contains the runtime shape. |
| 346 | if self.compile_sizes is None: |
| 347 | return None |
| 348 | |
| 349 | if runtime_shape in self.compile_sizes: |
| 350 | return self.range_entries[Range(start=runtime_shape, end=runtime_shape)] |
| 351 | else: |
| 352 | for range in self.compile_ranges: |
| 353 | if runtime_shape in range: |
| 354 | return self.range_entries[range] |
| 355 | return None |
| 356 | |
| 357 | def __call__(self, *args: Any) -> Any: |
| 358 | if self.sym_shape_indices: |
| 359 | runtime_shape = args[self.sym_shape_indices[0]] |
| 360 | range_entry = self._find_range_for_shape(runtime_shape) |
| 361 | assert range_entry is not None, ( |
| 362 | f"Shape: {runtime_shape} out of considered ranges: " |
| 363 | f"{self.compile_ranges}" |
| 364 | ) |
| 365 | else: |
| 366 | # All inputs have static shapes; use the only compiled range_entry |
| 367 | compiled_entries = [re for re in self.range_entries.values() if re.compiled] |
| 368 | assert len(compiled_entries) == 1, ( |
| 369 | f"Expected exactly one compiled range_entry for static shape " |
| 370 | f"compilation, but found {len(compiled_entries)}" |
| 371 | ) |
| 372 | range_entry = compiled_entries[0] |
| 373 | |
| 374 | assert range_entry.compiled, ( |
| 375 | "All ranges should be compiled or loaded up front in " |
| 376 | "PiecewiseBackend.__init__. " |
| 377 | f"range_entry={range_entry.compile_range}" |
| 378 | ) |
| 379 | return range_entry.runnable(*args) |