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 6, 15:45 UTC
[compile] Enable mega aot artifact for torch 2.12+. (#37198)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Mar 16, 2026 at 17:05 UTC
e6ae4b1be1c3dca1c25d7a12058dbb1fd900caa2
2 files changed
+15
-12
vllm/compilation/caching.py
+4
-8
@@ -307,13 +307,6 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
307
num_submods = len(submod_names)
308
num_artifacts = standalone_compile_artifacts.num_artifacts()
309
310
- logger.info(
311
- "reconstructing serializable fn from standalone compile "
312
- "artifacts. num_artifacts=%d num_submods=%d",
313
- num_artifacts,
314
- num_submods,
315
- )
316
-
310
with functorch_ctx:
311
fn = reconstruct_serializable_fn_from_mega_artifact(
312
state=state,
@@ -324,7 +317,10 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
317
)
318
319
logger.info(
327
- "reconstructed serializable fn from standalone compile artifacts"
320
+ "reconstructed serializable fn from standalone compile "
321
+ "artifacts. num_artifacts=%d num_submods=%d",
322
+ num_artifacts,
323
+ num_submods,
324
)
325
326
return fn
vllm/envs.py
+11
-4
@@ -296,6 +296,16 @@ def use_aot_compile() -> bool:
296
)
297
298
299
+def use_mega_aot_artifact():
300
+ from vllm.utils.torch_utils import is_torch_equal_or_newer
301
+
302
+ default_value = (
303
+ "1" if is_torch_equal_or_newer("2.12.0.dev") and use_aot_compile() else "0"
304
+ )
305
+
306
+ return os.environ.get("VLLM_USE_MEGA_AOT_ARTIFACT", default_value) == "1"
307
+
308
+
309
def env_with_choices(
310
env_name: str,
311
default: str | None,
@@ -616,10 +626,7 @@ environment_variables: dict[str, Callable[[], Any]] = {
626
# Enable loading compiled models directly from cached standalone compile artifacts
627
# without re-splitting graph modules. This reduces overhead during model
628
# loading by using reconstruct_serializable_fn_from_mega_artifact.
619
- "VLLM_USE_MEGA_AOT_ARTIFACT": lambda: os.environ.get(
620
- "VLLM_USE_MEGA_AOT_ARTIFACT", "0"
621
- )
622
- == "1",
629
+ "VLLM_USE_MEGA_AOT_ARTIFACT": use_mega_aot_artifact,
630
# local rank of the process in the distributed setting, used to determine
631
# the GPU device id
632
"LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),