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

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 15, 06:00 UTC

[torch.compile] reorganize the cache directory to support compiling multiple models (#19064)

Signed-off-by: youkaichao <youkaichao@gmail.com>

youkaichao committed Jun 13, 2025 at 15:23 UTC d70bc7c02957ff6c11bc117527f7c724bb6aeb46

6 files changed +117 -27

vllm/compilation/backends.py

+56 -9

@@ -7,6 +7,7 @@ import os

7 import pprint

8 import time

9 from collections.abc import Sequence

10 +from contextlib import contextmanager

11 from typing import Any, Callable, Optional

13 import torch

@@ -66,7 +67,25 @@ class CompilerManager:

67 def compute_hash(self, vllm_config: VllmConfig) -> str:

68 return self.compiler.compute_hash(vllm_config)

69 - def initialize_cache(self, cache_dir: str, disable_cache: bool = False):

70 + def initialize_cache(self,

71 + cache_dir: str,

72 + disable_cache: bool = False,

73 + prefix: str = ""):

74 + """

75 + Initialize the cache directory for the compiler.

76 +

77 + The organization of the cache directory is as follows:

78 + cache_dir=/path/to/hash_str/rank_i_j/prefix/

79 + inside cache_dir, there will be:

80 + - vllm_compile_cache.py

81 + - computation_graph.py

82 + - transformed_code.py

83 +

84 + for multiple prefixes, they can share the same

85 + base cache dir of /path/to/hash_str/rank_i_j/ ,

86 + to store some common compilation artifacts.

87 + """

88 +

89 self.disable_cache = disable_cache

90 self.cache_dir = cache_dir

91 self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")

@@ -80,7 +99,8 @@ class CompilerManager:

99 self.cache = ast.literal_eval(f.read())

100

101 self.compiler.initialize_cache(cache_dir=cache_dir,

83 - disable_cache=disable_cache)

102 + disable_cache=disable_cache,

103 + prefix=prefix)

104

105 def save_to_file(self):

106 if self.disable_cache or not self.is_cache_updated:

@@ -310,6 +330,25 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):

330 return output

331

332

333 +# the tag for the part of model being compiled,

334 +# e.g. backbone/eagle_head

335 +model_tag: str = "backbone"

336 +

337 +

338 +@contextmanager

339 +def set_model_tag(tag: str):

340 + """Context manager to set the model tag."""

341 + global model_tag

342 + assert tag != model_tag, \

343 + f"Model tag {tag} is the same as the current tag {model_tag}."

344 + old_tag = model_tag

345 + model_tag = tag

346 + try:

347 + yield

348 + finally:

349 + model_tag = old_tag

350 +

351 +

352 class VllmBackend:

353 """The compilation backend for `torch.compile` with vLLM.

354 It is used for compilation level of `CompilationLevel.PIECEWISE`,

@@ -341,7 +380,17 @@ class VllmBackend:

380 def __init__(

381 self,

382 vllm_config: VllmConfig,

383 + prefix: str = "",

384 ):

385 +

386 + # if the model is initialized with a non-empty prefix,

387 + # then usually it's enough to use that prefix,

388 + # e.g. launguage_model, vision_model, etc.

389 + # when multiple parts are initialized as independent

390 + # models, we need to use the model_tag to distinguish

391 + # them, e.g. backbone (default), eagle_head, etc.

392 + self.prefix = prefix or model_tag

393 +

394 global global_graph_pool

395 if global_graph_pool is None:

396 global_graph_pool = current_platform.graph_pool_handle()

@@ -441,16 +490,13 @@ class VllmBackend:

490 )

491 self.compilation_config.cache_dir = cache_dir

492

444 - if compilation_counter.num_graphs_seen > 0:

445 - cache_dir = self.compilation_config.cache_dir + \

446 - f'-{compilation_counter.num_graphs_seen}'

447 - else:

448 - cache_dir = self.compilation_config.cache_dir

493 + cache_dir = self.compilation_config.cache_dir

494 os.makedirs(cache_dir, exist_ok=True)

495 self.compilation_config.cache_dir = cache_dir

496 rank = vllm_config.parallel_config.rank

497 dp_rank = vllm_config.parallel_config.data_parallel_rank

453 - local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")

498 + local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}",

499 + self.prefix)

500 os.makedirs(local_cache_dir, exist_ok=True)

501 self.compilation_config.local_cache_dir = local_cache_dir

502

@@ -462,7 +508,8 @@ class VllmBackend:

508 logger.info("Using cache directory: %s for vLLM's torch.compile",

509 local_cache_dir)

510

465 - self.compiler_manager.initialize_cache(local_cache_dir, disable_cache)

511 + self.compiler_manager.initialize_cache(local_cache_dir, disable_cache,

512 + self.prefix)

513

514 # when dynamo calls the backend, it means the bytecode

515 # transform and analysis are done

vllm/compilation/compiler_interface.py

+29 -9

@@ -28,11 +28,22 @@ class CompilerInterface:

28 # This is a class-level attribute.

29 name: str

31 - def initialize_cache(self, cache_dir: str, disable_cache: bool = False):

31 + def initialize_cache(self,

32 + cache_dir: str,

33 + disable_cache: bool = False,

34 + prefix: str = ""):

35 """

36 when the vLLM process uses `cache_dir` as the cache directory,

37 the compiler should initialize itself with the cache directory,

38 e.g. by re-directing its own cache directory to a sub-directory.

39 +

40 + prefix can be used in combination with cache_dir to figure out the base

41 + cache directory, e.g. there're multiple parts of model being compiled,

42 + but we want to share the same cache directory for all of them.

43 +

44 + e.g.

45 + cache_dir = "/path/to/dir/backbone", prefix = "backbone"

46 + cache_dir = "/path/to/dir/eagle_head", prefix = "eagle_head"

47 """

48 pass

@@ -166,7 +177,10 @@ class InductorStandaloneAdaptor(CompilerInterface):

177 usedforsecurity=False).hexdigest()[:10]

178 return hash_str

179

169 - def initialize_cache(self, cache_dir: str, disable_cache: bool = False):

180 + def initialize_cache(self,

181 + cache_dir: str,

182 + disable_cache: bool = False,

183 + prefix: str = ""):

184 self.cache_dir = cache_dir

185

186 def compile(

@@ -242,18 +256,23 @@ class InductorAdaptor(CompilerInterface):

256 usedforsecurity=False).hexdigest()[:10]

257 return hash_str

258

245 - def initialize_cache(self, cache_dir: str, disable_cache: bool = False):

259 + def initialize_cache(self,

260 + cache_dir: str,

261 + disable_cache: bool = False,

262 + prefix: str = ""):

263 self.cache_dir = cache_dir

264 + self.prefix = prefix

265 + self.base_cache_dir = cache_dir[:-len(prefix)] if prefix else cache_dir

266 if disable_cache:

267 return

268 # redirect the cache directory to a sub-directory

269 # set flags so that Inductor and Triton store their cache

270 # in the cache_dir, then users only need to copy the cache_dir

271 # to another machine to reuse the cache.

253 - inductor_cache = os.path.join(cache_dir, "inductor_cache")

272 + inductor_cache = os.path.join(self.base_cache_dir, "inductor_cache")

273 os.makedirs(inductor_cache, exist_ok=True)

274 os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache

256 - triton_cache = os.path.join(cache_dir, "triton_cache")

275 + triton_cache = os.path.join(self.base_cache_dir, "triton_cache")

276 os.makedirs(triton_cache, exist_ok=True)

277 os.environ["TRITON_CACHE_DIR"] = triton_cache

278

@@ -298,14 +317,14 @@ class InductorAdaptor(CompilerInterface):

317 nonlocal file_path

318 compiled_fn = inductor_compiled_graph.current_callable

319 file_path = compiled_fn.__code__.co_filename # noqa

301 - if not file_path.startswith(self.cache_dir):

320 + if not file_path.startswith(self.base_cache_dir):

321 # hooked in the align_inputs_from_check_idxs function

322 # in torch/_inductor/utils.py

323 for cell in compiled_fn.__closure__:

324 if not callable(cell.cell_contents):

325 continue

326 if cell.cell_contents.__code__.co_filename.startswith(

308 - self.cache_dir):

327 + self.base_cache_dir):

328 # this is the real file path compiled from Inductor

329 file_path = cell.cell_contents.__code__.co_filename

330 break

@@ -325,14 +344,15 @@ class InductorAdaptor(CompilerInterface):

344 nonlocal file_path

345 compiled_fn = inductor_compiled_graph.current_callable

346 file_path = compiled_fn.__code__.co_filename # noqa

328 - if not file_path.startswith(self.cache_dir):

347 + if not file_path.startswith(self.base_cache_dir):

348 # hooked in the align_inputs_from_check_idxs function

349 # in torch/_inductor/utils.py

350 for cell in compiled_fn.__closure__:

351 if not callable(cell.cell_contents):

352 continue

353 code = cell.cell_contents.__code__

335 - if code.co_filename.startswith(self.cache_dir):

354 + if code.co_filename.startswith(

355 + self.base_cache_dir):

356 # this is the real file path

357 # compiled from Inductor

358 file_path = code.co_filename

vllm/config.py

+17 -2

@@ -4666,10 +4666,13 @@ class VllmConfig:

4666

4667

4668 _current_vllm_config: Optional[VllmConfig] = None

4669 +_current_prefix: Optional[str] = None

4670

4671

4672 @contextmanager

4672 -def set_current_vllm_config(vllm_config: VllmConfig, check_compile=False):

4673 +def set_current_vllm_config(vllm_config: VllmConfig,

4674 + check_compile=False,

4675 + prefix: Optional[str] = None):

4676 """

4677 Temporarily set the current vLLM config.

4678 Used during model initialization.

@@ -4677,12 +4680,14 @@ def set_current_vllm_config(vllm_config: VllmConfig, check_compile=False):

4680 so that all modules can access it, e.g. custom ops

4681 can access the vLLM config to determine how to dispatch.

4682 """

4680 - global _current_vllm_config

4683 + global _current_vllm_config, _current_prefix

4684 old_vllm_config = _current_vllm_config

4685 + old_prefix = _current_prefix

4686 from vllm.compilation.counter import compilation_counter

4687 num_models_seen = compilation_counter.num_models_seen

4688 try:

4689 _current_vllm_config = vllm_config

4690 + _current_prefix = prefix

4691 yield

4692 except Exception:

4693 raise

@@ -4706,6 +4711,7 @@ def set_current_vllm_config(vllm_config: VllmConfig, check_compile=False):

4711 vllm_config.model_config.model)

4712 finally:

4713 _current_vllm_config = old_vllm_config

4714 + _current_prefix = old_prefix

4715

4716

4717 def get_current_vllm_config() -> VllmConfig:

@@ -4719,6 +4725,15 @@ def get_current_vllm_config() -> VllmConfig:

4725 return _current_vllm_config

4726

4727

4728 +def get_current_model_prefix() -> str:

4729 + """

4730 + Get the prefix of the model that's currently being initialized.

4731 + """

4732 + assert _current_prefix is not None, \

4733 + "Current model prefix is not set. "

4734 + return _current_prefix

4735 +

4736 +

4737 def contains_object_print(text):

4738 """

4739 Check if the text looks like a printed Python object, e.g.

vllm/model_executor/model_loader/utils.py

+6 -2

@@ -58,7 +58,9 @@ def initialize_model(

58 all_params = [param.name for param in signatures.parameters.values()]

59 if "vllm_config" in all_params and "prefix" in all_params:

60 # new-style model class

61 - with set_current_vllm_config(vllm_config, check_compile=True):

61 + with set_current_vllm_config(vllm_config,

62 + check_compile=True,

63 + prefix=prefix):

64 return model_class(vllm_config=vllm_config, prefix=prefix)

66 msg = ("vLLM model class should accept `vllm_config` and `prefix` as "

@@ -86,7 +88,9 @@ def initialize_model(

88 kwargs["lora_config"] = vllm_config.lora_config

89 if "scheduler_config" in all_params:

90 kwargs["scheduler_config"] = vllm_config.scheduler_config

89 - with set_current_vllm_config(vllm_config, check_compile=True):

91 + with set_current_vllm_config(vllm_config,

92 + check_compile=True,

93 + prefix=prefix):

94 return model_class(**kwargs)

vllm/v1/spec_decode/eagle.py

+4 -2

@@ -320,8 +320,10 @@ class EagleProposer:

320 target_attn_layer_names = set(

321 get_layers_from_vllm_config(self.vllm_config, Attention).keys())

322

323 - self.model = get_model(vllm_config=self.vllm_config,

324 - model_config=draft_model_config)

323 + from vllm.compilation.backends import set_model_tag

324 + with set_model_tag("eagle_head"):

325 + self.model = get_model(vllm_config=self.vllm_config,

326 + model_config=draft_model_config)

327

328 draft_attn_layer_names = (

329 get_layers_from_vllm_config(self.vllm_config, Attention).keys() -

vllm/v1/spec_decode/medusa.py

+5 -3

@@ -48,9 +48,11 @@ class MedusaProposer:

48 return [list(row) for row in zip(*draft_tokens)]

50 def load_model(self, target_model: nn.Module) -> None:

51 - self.model = get_model(vllm_config=self.vllm_config,

52 - model_config=self.vllm_config.

53 - speculative_config.draft_model_config)

51 + from vllm.compilation.backends import set_model_tag

52 + with set_model_tag("medusa_head"):

53 + self.model = get_model(vllm_config=self.vllm_config,

54 + model_config=self.vllm_config.

55 + speculative_config.draft_model_config)

57 @torch.inference_mode()

58 def dummy_run(self, num_tokens: int) -> None: