@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5569f5218d

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 7, 06:30 UTC

[torch.compile] Stop lazily compiling (#35472)

Signed-off-by: Richard Zou <zou3519@gmail.com>

Richard Zou committed Mar 4, 2026 at 15:13 UTC 5569f5218d3b8a08cfbb9fd51c9f01852f16ddbc

7 files changed +177 -150

tests/compile/test_compile_ranges.py

+4 -3

@@ -73,6 +73,7 @@ def test_compile_ranges(use_fresh_inductor_cache):

73 Range(start=16, end=16),

74 Range(start=9, end=32),

75 Range(start=64, end=64),

76 + Range(start=128, end=128),

77 Range(start=33, end=8192),

78 ]

79 )

@@ -95,16 +96,16 @@ def test_compile_ranges(use_fresh_inductor_cache):

97 with set_current_vllm_config(vllm_config):

98 model = TestModel(vllm_config=vllm_config, prefix="").eval()

98 - # Number of compilations: 3 for each compile range + 2 compile sizes

99 + # Number of compilations: 3 compile ranges + 3 compile sizes

100 batch_sizes = [1, 4, 16, 24, 48, 64, 8192]

101

102 with compilation_counter.expect(

103 num_graphs_seen=1,

104 num_piecewise_graphs_seen=1,

104 - num_backend_compilations=5,

105 + num_backend_compilations=6,

106 ):

107 run_model(vllm_config, model, batch_sizes)

107 - assert post_grad_range_checker.num_calls == 5

108 + assert post_grad_range_checker.num_calls == 6

109

110

111 def test_compile_config_get_compile_ranges():

tests/compile/test_structured_logging.py

+3 -3

@@ -109,9 +109,9 @@ def test_vllm_structured_logging_artifacts(use_fresh_inductor_cache):

109 f"got {len(vllm_piecewise_split_graph)}"

110 )

111 compile_start_artifacts = capture.get("artifact", "vllm_piecewise_compile_start")

112 - assert len(compile_start_artifacts) == 2, (

113 - "Expected 2 vllm_piecewise_compile_start "

114 - "(one for dynamic ranges, one for compile size), "

112 + assert len(compile_start_artifacts) == 4, (

113 + "Expected 4 vllm_piecewise_compile_start "

114 + "(2 subgraphs x 2 ranges each: dynamic + compile size), "

115 f"got {len(compile_start_artifacts)}"

116 )

117 submod_dumps = capture.get("graph_dump", r"vllm_submod_.*")

vllm/compilation/backends.py

+20 -31

@@ -2,7 +2,6 @@

2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project

4 import ast

5 -import contextvars

5 import dataclasses

6 import hashlib

7 import json

@@ -18,7 +17,7 @@ from typing import Any

18 import torch

19 import torch.fx as fx

21 -from torch._dispatch.python import enable_python_dispatcher

20 +from torch._dynamo.utils import dynamo_timed

21 from torch._logging._internal import trace_structured

23 import vllm.envs as envs

@@ -510,9 +509,9 @@ def wrap_with_cudagraph_if_needed(

509

510 class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]

511 """Code adapted from `torch.fx.passes.shape_prop.ShapeProp`.

513 - It runs the given graph with fake inputs, and compile some

514 - submodules specified by `compile_submod_names` with the given

515 - compilation configs.

512 + It runs the given split graph interpreter, and for each submodule in

513 + `compile_submod_names`, creates a PiecewiseBackend and compiles all

514 + ranges up front.

515

516 NOTE: the order in `compile_submod_names` matters, because

517 it will be used to determine the order of the compiled piecewise

@@ -540,9 +539,6 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]

539 vllm_backend: "VllmBackend",

540 ) -> None:

541 super().__init__(module)

543 - from torch._guards import detect_fake_mode

544 -

545 - self.fake_mode = detect_fake_mode()

542 self.compile_submod_names = compile_submod_names

543 self.compilation_config = vllm_config.compilation_config

544 self.vllm_config = vllm_config

@@ -552,13 +548,7 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]

548

549 @instrument(span_name="Inductor compilation")

550 def run(self, *args: Any) -> Any:

555 - # maybe instead just assert inputs are fake?

556 - fake_args = [

557 - self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t

558 - for t in args

559 - ]

560 - with self.fake_mode, enable_python_dispatcher():

561 - return super().run(*fake_args)

551 + return super().run(*args)

552

553 def call_module(

554 self,

@@ -614,21 +604,6 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]

604 model_tag: str = "backbone"

605 model_is_encoder: bool = False

606

617 -_on_compilation_complete_callback: contextvars.ContextVar[Callable[[], None] | None] = (

618 - contextvars.ContextVar("on_compilation_complete_callback", default=None)

619 -)

620 -

621 -

622 -@contextmanager

623 -def set_on_compilation_complete(

624 - callback: Callable[[], None],

625 -) -> Generator[None, None, None]:

626 - token = _on_compilation_complete_callback.set(callback)

627 - try:

628 - yield

629 - finally:

630 - _on_compilation_complete_callback.reset(token)

631 -

607

608 @contextmanager

609 def set_model_tag(tag: str, is_encoder: bool = False) -> Generator[None, None, None]:

@@ -846,6 +821,7 @@ class VllmBackend:

821 ),

822 )

823

824 + @dynamo_timed("vllm_backend")

825 def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any:

826 from .caching import (

827 VllmSerializableFunction,

@@ -1036,11 +1012,24 @@ class VllmBackend:

1012 ]

1013

1014 # propagate the split graph to the piecewise backend,

1039 - # compile submodules with symbolic shapes

1015 + # compile submodules with symbolic shapes, and compile all ranges

1016 + # up front so that compilation is complete before the callable

1017 + # is returned.

1018 PiecewiseCompileInterpreter(

1019 self.split_gm, submod_names_to_compile, self.vllm_config, self

1020 ).run(*fake_args)

1021

1022 + # All compilation is done. Save the cache.

1023 + time_before_saving = time.perf_counter()

1024 + self.compiler_manager.save_to_file()

1025 + elapsed = time.perf_counter() - time_before_saving

1026 + if elapsed > 1:

1027 + logger.info_once(

1028 + "Saved compiler manager cache in %.2f seconds.",

1029 + elapsed,

1030 + scope="local",

1031 + )

1032 +

1033 from torch._guards import detect_fake_mode

1034

1035 fake_mode = detect_fake_mode()

vllm/compilation/caching.py

+11 -15

@@ -313,30 +313,26 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]

313

314 return fn

315

316 - # Fall back to standard VllmBackend

316 + # Fall back to standard VllmBackend.

317 + # Use a lazy closure: the backend needs traced_files for cache

318 + # dir computation, but those are only populated after

319 + # _verify_source_unchanged runs in decorators.py (which happens

320 + # after deserialization completes).

321 from vllm.compilation.backends import VllmBackend

322

323 is_encoder = state.get("is_encoder", False)

320 - vllm_backend: VllmBackend = VllmBackend(

321 - get_current_vllm_config(), state["prefix"], is_encoder

322 - )

324 + vllm_config = get_current_vllm_config()

325 + compile_inputs = list(state["example_inputs"])

326

327 def optimized_call(*example_inputs: Any) -> Any:

325 - """

326 - On the first run of the optimized call, we rerun the compiler

327 - backend which should result in a cache hit. After the backend

328 - call returns, we just do a one-time replacement of the optimized

329 - call with the compiled function, so that subsequent calls are on

330 - the AOT compiled path.

331 - """

332 - compile_inputs = [

333 - inp if inp is not None else example_inputs[i]

334 - for i, inp in enumerate(fn.example_inputs)

335 - ]

328 + vllm_backend: VllmBackend = VllmBackend(

329 + vllm_config, state["prefix"], is_encoder

330 + )

331 with tracing(TracingContext(fake_mode)):

332 fn.optimized_call = vllm_backend(

333 state["graph_module"], compile_inputs

334 ).optimized_call

335 + fn.vllm_backend = vllm_backend

336 return fn.optimized_call(*example_inputs)

337

338 fn = cls(**state, optimized_call=optimized_call)

vllm/compilation/decorators.py

+12 -7

@@ -466,8 +466,12 @@ def _support_torch_compile(

466 "Directly load AOT compilation from path %s", aot_compilation_path

467 )

468 # Apply partition wrapper context for proper CUDA graph capture

469 + from .monitor import end_monitoring_torch_compile

470 +

471 with maybe_use_cudagraph_partition_wrapper(self.vllm_config):

470 - return self.aot_compiled_fn(self, *args, **kwargs)

472 + output = self.aot_compiled_fn(self, *args, **kwargs)

473 + end_monitoring_torch_compile(self.vllm_config)

474 + return output

475

476 if self.compiled:

477 assert (

@@ -552,18 +556,19 @@ def _support_torch_compile(

556 logger.warning("Detected eager backend, disabling AOT compile.")

557 use_aot_compile = False

558 if use_aot_compile:

555 - from vllm.compilation.backends import set_on_compilation_complete

556 -

559 # store the path for saving after warmup

560 self._aot_compilation_path = aot_compilation_path

561 self._aot_cache_dir = cache_dir

560 - # set callback in context so it's available when compilation completes

561 - with set_on_compilation_complete(self.save_aot_compiled_function):

562 - self.aot_compiled_fn = self.aot_compile(*args, **kwargs)

563 - output = self.aot_compiled_fn(self, *args, **kwargs)

562 + self.aot_compiled_fn = self.aot_compile(*args, **kwargs)

563 + # All compilation is done at this point, save the AOT artifact.

564 + self.save_aot_compiled_function()

565 + output = self.aot_compiled_fn(self, *args, **kwargs)

566 else:

567 output = TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs) # type: ignore[arg-type]

568

569 + from .monitor import end_monitoring_torch_compile

570 +

571 + end_monitoring_torch_compile(self.vllm_config)

572 self.compiled = True

573 return output

574

vllm/compilation/monitor.py

+1 -1

@@ -33,7 +33,7 @@ def end_monitoring_torch_compile(vllm_config: VllmConfig) -> None:

33 total_compile_time: float = time.perf_counter() - torch_compile_start_time

34 if compilation_config.mode == CompilationMode.VLLM_COMPILE:

35 logger.info_once(

36 - "torch.compile takes %.2f s in total",

36 + "torch.compile and initial profiling run took %.2f s in total",

37 total_compile_time,

38 scope="local",

39 )

vllm/compilation/piecewise_backend.py

+126 -90

@@ -5,7 +5,6 @@ import dataclasses

5 import io

6 import json

7 import pickle

8 -import time

8 from collections.abc import Callable

9 from pickle import Pickler

10 from typing import Any

@@ -16,7 +15,6 @@ from torch._inductor.runtime.triton_heuristics import CachingAutotuner

15 from torch._logging._internal import trace_structured

17 from vllm.compilation.backends import VllmBackend

19 -from vllm.compilation.monitor import end_monitoring_torch_compile

18 from vllm.config import VllmConfig

19 from vllm.config.utils import Range

20 from vllm.logger import init_logger

@@ -24,6 +22,55 @@ from vllm.logger import init_logger

22 logger = init_logger(__name__)

25 +def get_fake_args_from_graph(graph: fx.GraphModule) -> list[Any]:

26 + """Get fake args directly from graph placeholder nodes."""

27 + fake_args = []

28 + for node in graph.graph.nodes:

29 + if node.op == "placeholder":

30 + fake_args.append(node.meta["example_value"])

31 + else:

32 + break

33 + return fake_args

34 +

35 +

36 +def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]:

37 + """Create example inputs with symbolic dims replaced by a concrete size.

38 +

39 + Used for single-size eager compilation where we need concrete-shaped

40 + inputs but don't have real runtime tensors yet.

41 + """

42 + from torch._prims_common import compute_required_storage_length

43 + from torch.fx.experimental.symbolic_shapes import is_symbolic

44 +

45 + def concretize(sym_val: Any) -> int:

46 + """Replace all symbolic variables in a SymInt expression with size."""

47 + if not is_symbolic(sym_val):

48 + return int(sym_val)

49 + expr = sym_val.node.expr

50 + return int(expr.subs({s: size for s in expr.free_symbols}))

51 +

52 + args: list[Any] = []

53 + for node in graph.graph.nodes:

54 + if node.op != "placeholder":

55 + break

56 + val = node.meta["example_value"]

57 + if isinstance(val, torch.SymInt):

58 + args.append(concretize(val))

59 + elif isinstance(val, torch.Tensor):

60 + new_shape = tuple(concretize(d) for d in val.shape)

61 + new_strides = tuple(concretize(s) for s in val.stride())

62 + new_storage_offset = concretize(val.storage_offset())

63 + needed_size = compute_required_storage_length(

64 + new_shape, new_strides, new_storage_offset

65 + )

66 + t = torch.empty(needed_size, dtype=val.dtype, device=val.device)

67 + t = t.as_strided(new_shape, new_strides, new_storage_offset)

68 + args.append(t)

69 + else:

70 + args.append(val)

71 + return args

72 +

73 +

74 @dataclasses.dataclass

75 class RangeEntry:

76 compile_range: Range

@@ -109,10 +156,6 @@ class PiecewiseBackend:

156 # the entries for ranges that we need to either

157 self.range_entries: dict[Range, RangeEntry] = {}

158

112 - # to_be_compiled_ranges tracks the remaining ranges to compile,

113 - # and updates during the compilation process, so we need to copy it

114 - self.to_be_compiled_ranges: set[Range] = set(self.compile_ranges)

115 -

159 # We only keep compilation management inside this class directly.

160 if self.compile_sizes is not None:

161 for size in self.compile_sizes:

@@ -129,7 +172,6 @@ class PiecewiseBackend:

172 self.range_entries[range] = RangeEntry(

173 compile_range=range,

174 )

132 - self.to_be_compiled_ranges.add(range)

175

176 for range in self.compile_ranges:

177 self.range_entries[range] = RangeEntry(

@@ -139,12 +181,10 @@ class PiecewiseBackend:

181 # Track whether we've logged the graph for this subgraph (only log once)

182 self._graph_logged = False

183

142 - # get the on_compilation_complete callback from context...

143 - # PiecewiseBackend is created during the first call,

144 - # which is when the context is set (see compilation/decorators.py)

145 - from vllm.compilation.backends import _on_compilation_complete_callback

146 -

147 - self.on_compilation_complete = _on_compilation_complete_callback.get()

184 + if self.graph is not None:

185 + self.compile_all_ranges()

186 + else:

187 + self.load_all_ranges()

188

189 def get_compiled_graph_wrapper(

190 self, compiled_graph: Callable[..., Any]

@@ -161,25 +201,6 @@ class PiecewiseBackend:

201

202 return compiled_graph_wrapper

203

164 - def check_for_ending_compilation(self) -> None:

165 - if self.is_last_graph and not self.to_be_compiled_ranges:

166 - # no specific sizes to compile

167 - # save the hash of the inductor graph for the next run

168 - time_before_saving = time.perf_counter()

169 - self.vllm_backend.compiler_manager.save_to_file()

170 - elapsed = time.perf_counter() - time_before_saving

171 - if elapsed > 1:

172 - logger.info_once(

173 - "Saved compiler manager cache in %.2f seconds.",

174 - elapsed,

175 - scope="local",

176 - )

177 -

178 - end_monitoring_torch_compile(self.vllm_config)

179 - # Call the completion callback (e.g., to save AOT compiled function)

180 - if self.on_compilation_complete is not None:

181 - self.on_compilation_complete()

182 -

204 def to_bytes(self) -> dict[str, bytes]:

205 class StandaloneCompiledArtifactsPickler(Pickler):

206 def reducer_override(self, obj: object) -> Any:

@@ -216,27 +237,54 @@ class PiecewiseBackend:

237

238 return out

239

219 - def _fakify_args(self, args: tuple[Any, ...]) -> list[Any]:

220 - # We need to pass fake example_inputs, otherwise torch.compile

221 - # will fakify the example_inputs potentially causing some non dynamic

222 - # dimension to be be duck shaped to other existing shapes that have hints

223 - # matching their values.

224 - # This is problem because it can lead to unintended specializations!

225 - # if the new wrongly dynamic dim is specialized

226 - # it will force specializing the whole shape

227 - # torch.compile probably should not accept

228 - # non fake tensors as example inputs!

229 - # See issue https://github.com/vllm-project/vllm/issues/27899

230 - fake_example_inputs = []

231 - assert self.graph is not None

232 - for node in self.graph.graph.nodes:

233 - # All place holders come first

234 - if node.op == "placeholder":

235 - fake_example_inputs.append(node.meta["example_value"])

240 + def compile_all_ranges(self) -> None:

241 + """Compile all range entries for this piecewise subgraph up front."""

242 + assert self.graph is not None, (

243 + "Cannot compile without a graph. "

244 + "When loading from cache/AOT artifacts, "

245 + "compile_all_ranges should not be called."

246 + )

247 +

248 + for range_entry in self.range_entries.values():

249 + if range_entry.compiled:

250 + continue

251 +

252 + self._log_compile_start(range_entry.compile_range)

253 +

254 + if range_entry.compile_range.is_single_size():

255 + args_list = create_concrete_args(

256 + self.graph, range_entry.compile_range.start

257 + )

258 else:

237 - break

238 - assert len(fake_example_inputs) == len(args)

239 - return fake_example_inputs

259 + args_list = get_fake_args_from_graph(self.graph)

260 +

261 + # TODO(https://github.com/vllm-project/vllm/issues/35766)

262 + # Can we remove strict_autograd_cache and

263 + # force_non_lazy_backward_lowering overrides?

264 + # I added them explicitly because this is what they are

265 + # set to before the refactor

266 + # (https://github.com/vllm-project/vllm/pull/35472).

267 + # They affect the aotautograd cache key computation

268 + # but they shouldn't have any effect on the actual

269 + # compilation.

270 + config_patches = dict(

271 + bundled_autograd_cache=True,

272 + strict_autograd_cache=False,

273 + )

274 + if hasattr(torch._functorch.config, "force_non_lazy_backward_lowering"):

275 + config_patches["force_non_lazy_backward_lowering"] = False

276 + with torch._functorch.config.patch(**config_patches):

277 + range_entry.runnable = self.vllm_backend.compiler_manager.compile(

278 + self.graph,

279 + args_list,

280 + self.vllm_backend.inductor_config,

281 + self.compilation_config,

282 + compile_range=range_entry.compile_range,

283 + graph_index=self.piecewise_compile_index,

284 + num_graphs=self.total_piecewise_compiles,

285 + )

286 +

287 + range_entry.compiled = True

288

289 def _log_compile_start(self, compile_range: Range):

290 """Log compilation event for TORCH_TRACE/tlparse."""

@@ -277,44 +325,29 @@ class PiecewiseBackend:

325 payload_fn=lambda: self.graph.print_readable(print_output=False),

326 )

327

280 - def _maybe_compile_for_range_entry(

281 - self, range_entry: RangeEntry, args: tuple[Any, ...]

282 - ) -> Any:

283 - if not range_entry.compiled:

284 - if self.compiled_runnables is not None:

285 - range_entry.runnable = self.get_compiled_graph_wrapper(

286 - self.compiled_runnables[str(range_entry.compile_range)]

287 - )

288 - else:

289 - self._log_compile_start(range_entry.compile_range)

290 -

291 - # args are real arguments

292 - # fakify for range, real args for concrete size.

293 - # For concrete size, we clear the shape env in

294 - # compiler_manager.compile() so no need to fakify.

295 - args_list = (

296 - self._fakify_args(args)

297 - if not range_entry.compile_range.is_single_size()

298 - else list(args)

299 - )

300 -

301 - with (

302 - torch._functorch.config.patch("bundled_autograd_cache", True),

303 - ):

304 - range_entry.runnable = self.vllm_backend.compiler_manager.compile(

305 - self.graph,

306 - args_list,

307 - self.vllm_backend.inductor_config,

308 - self.compilation_config,

309 - compile_range=range_entry.compile_range,

310 - graph_index=self.piecewise_compile_index,

311 - num_graphs=self.total_piecewise_compiles,

312 - )

328 + def load_all_ranges(self) -> None:

329 + """Load all pre-compiled runnables for this piecewise subgraph.

330

331 + Called during warm start to wrap all cached compiled_runnables

332 + into range_entry.runnable up front, analogous to compile_all_ranges()

333 + for the cold start path.

334 + """

335 + assert self.compiled_runnables is not None, (

336 + "load_all_ranges should only be called when compiled_runnables "

337 + "is set (warm start / cache loading path)."

338 + )

339 + for range_entry in self.range_entries.values():

340 + if range_entry.compiled:

341 + continue

342 + key = str(range_entry.compile_range)

343 + assert key in self.compiled_runnables, (

344 + f"Missing compiled runnable for range {range_entry.compile_range}. "

345 + f"Available keys: {list(self.compiled_runnables.keys())}"

346 + )

347 + range_entry.runnable = self.get_compiled_graph_wrapper(

348 + self.compiled_runnables[key]

349 + )

350 range_entry.compiled = True

315 - self.to_be_compiled_ranges.remove(range_entry.compile_range)

316 -

317 - self.check_for_ending_compilation()

351

352 def _find_range_for_shape(self, runtime_shape: int) -> RangeEntry | None:

353 # First we try to find the range entry for the concrete compile size

@@ -338,6 +371,9 @@ class PiecewiseBackend:

371 assert range_entry is not None, (

372 f"Shape: {runtime_shape} out of considered ranges: {self.compile_ranges}"

373 )

341 -

342 - self._maybe_compile_for_range_entry(range_entry, args)

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)