[AOT compilation] support torch.compile inductor artifacts in VllmCompiledFunction by dolpm · Pull Request #25205 · vllm-project/vllm

dolpm · 2025-09-18T20:39:30Z

Adds support for inductor artifacts in VllmCompiledFunction

Purpose

we want to ensure that all pre-compiled artifacts (e.g., dynamo+inductor) live under one roof (i.e., VllmSerializableFunction).

the general implementation is as follows.

in the compile decorator, add support to only serialize/save the VllmCompiledFunction after all shapes ranges have been compiled. in the worst case, this will be done after the warmup phase when we encounter a cache miss, since this is when user-defined compile_sizes are run through the piecewise interpreter.
add a to_bytes method to cuda_piecewise_backend which will return a mapping from each shape to it's compiled bytes (using the function mentioned above).
when we serialize the VllmCompiledFunction, we iterate through all the compiled subgraphs and add these to a cache which can be serialized/deserialized and will dedupe the byte blobs. during deserialization, we will bulk-load these to get all of the runnables, and then re-map them to the submodule/shape pairs under which they were serialized originally.
when we load from the cache, we re-construct the split gm using the cached artifacts. we will pick the runnables from this cache directly.

Test Result

2026-01-12 17:29:23,876 - INFO - BENCHMARK SUMMARY
2026-01-12 17:29:23,876 - INFO - ========================================================================================================================
2026-01-12 17:29:23,876 - INFO - 
2026-01-12 17:29:23,876 - INFO - Model: Qwen/Qwen3-32B
2026-01-12 17:29:23,876 - INFO - --------------------------------------------------------------------------------
2026-01-12 17:29:23,876 - INFO -   Inductor: Cold=1m 21.44s, Warm=35.00s (2.33x speedup)
2026-01-12 17:29:23,876 - INFO -   Baseline: Cold=1m 10.50s, Warm=52.24s (1.35x speedup)
2026-01-12 17:29:23,876 - INFO -   Inductor vs Baseline: Cold=0.87x (-10.94s), Warm=1.49x (+17.25s)
2026-01-12 17:29:23,876 - INFO - 
2026-01-12 17:29:23,876 - INFO - Model: deepseek-ai/DeepSeek-R1-Distill-Qwen-32B
2026-01-12 17:29:23,876 - INFO - --------------------------------------------------------------------------------
2026-01-12 17:29:23,876 - INFO -   Inductor: Cold=1m 8.44s, Warm=36.10s (1.90x speedup)
2026-01-12 17:29:23,877 - INFO -   Baseline: Cold=1m 1.66s, Warm=49.95s (1.23x speedup)
2026-01-12 17:29:23,877 - INFO -   Inductor vs Baseline: Cold=0.90x (-6.78s), Warm=1.38x (+13.85s)
2026-01-12 17:29:23,877 - INFO - 
2026-01-12 17:29:23,877 - INFO - ==============================================================================================================

Essential Elements of an Effective PR Description Checklist

The purpose of the PR, such as "Fix some issue (link existing issues this PR will resolve)".
The test plan, such as providing test command.
The test results, such as pasting the results comparison before and after, or e2e results
(Optional) The necessary documentation update, such as updating supported_models.md and examples for a new model.
(Optional) Release notes update. If your change is user facing, please update the release notes draft in the Google Doc.

Note

Enables direct reuse of Inductor-compiled artifacts without re-splitting or recompiling, improving cold-start behavior and preserving output semantics.

Introduces VllmSerializableFunction.InductorCompiledArtifacts (deduped, pickleable, bulk-load) and embeds bytes plus sym_shape_indices/returns_tuple in serialized state
Adds VllmBackendWithCache to rebuild split_gm from cached artifacts and shared wrap_with_cudagraph_if_needed; validates cache completeness
Extends PiecewiseBackend with general-shape precompile for serialization, to_bytes(), preload from cache, and correct tuple vs tensor outputs; shared cudagraph wrapper and input copy helper
Updates VllmBackend to detect pre-split graphs, collect artifacts for serialization, expose on_compilation_complete, and pass sym-tensor indices for cudagraph copy
compiler_interface supports standalone_compile(aot=True) with fallback; includes inductor factors; adjusts functorch config
support_torch_compile: defer AOT save until post-warmup, add save_compile_cache, mark cache hits (aot_compile_loaded_from_cache)
Adds env VLLM_USE_BACKEND_WITH_INDUCTOR_COMPILED_ARTIFACTS, updates worker warmup to avoid CG padding, and extends pickle allowlist
New tests cover AOT save/load, tuple consistency, shape-env preservation, partition-wrapper on AOT load, GPT2 cache hit, and backend/cache units

^{Written by Cursor Bugbot for commit 3a4766931a62422c2c30e70b182f9f9dad1e99e7. This will update automatically on new commits. Configure here.}

mergify · 2025-09-18T20:40:19Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @dolpm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

mergify · 2025-09-25T18:32:37Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @dolpm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

zhxchen17

@dolpm Sorry for the delayed review. Just some initial thoughts.

mergify · 2025-10-24T18:48:11Z

This pull request has merge conflicts that must be resolved before it can be
merged. Please rebase the PR, @dolpm.

https://docs.github.com/en/pull-requests/collaborating-with-pull-requests/working-with-forks/syncing-a-fork

Summary: Fixing issue vllm-project#27348 For dynamo caching, it's possible that the compilation succeeds but the serialization step fails. In this case, the failure of serialization step shouldn't block user from getting compilation results correctly. Therefore we add a handling of the serialization error and only give warning when model saving fails. When saving fails, VLLM model runner should be able to just fallback to the old path, and in the next process, it will fail to load dynamo cache but still fallback to retracing with dynamo + loading inductor cache, which is the same behavior to AOT compile turned of off. This is mostly a short term fix and in the long term we should resolve the serialization bugs by eliminating pickling of graph modules. i.e. Once vllm-project#25205 is merged, we should be able to resolve the issue at a lower level. Test Plan: pytest tests/lora/test_quant_model.py Reviewers: Subscribers: Tasks: Tags: Signed-off-by: zhxchen17 <zhxchen17@fb.com>

zou3519 · 2026-01-14T22:45:16Z

NB: might merge conflict with #31756

zou3519

lgtm, let me know if you want to resolve any of the comments before merging

zou3519 · 2026-01-15T16:11:54Z

@@ -163,6 +467,11 @@ def compilation_config_hash_factors(vllm_config: VllmConfig) -> list[str]:
    #    model is created)
    config_hash = vllm_config.compute_hash()
    factors.append(config_hash)
+
+    # 2. inductor factors if applicable


There is an open problem here in that I think we need #29108 .

I think we're missing the "source code of the files that get traced through" in the hash. Zhengxu's work was going to make it so that vLLM includes "all python files" in the hash so we don't need to distinguish between which files get traced through.

basically, this stuff:

vllm/vllm/compilation/backends.py

Lines 588 to 627 in 130d6c9

env_factors = envs.compile_factors()

env_hash = hash_factors(env_factors)

# Compute config/compiler/code hashes once and reuse

config_hash = vllm_config.compute_hash()

compiler_hash = self.compiler_manager.compute_hash(vllm_config)

forward_code_files = list(sorted(self.compilation_config.traced_files))

logger.debug(

"Traced files (to be considered for compilation cache):\n%s",

lazy(lambda: "\n".join(forward_code_files)),

)

hash_content = []

for filepath in forward_code_files:

hash_content.append(filepath)

if filepath == "<string>":

# This means the function was dynamically generated, with

# e.g. exec(). We can't actually check these.

continue

try:

with open(filepath) as f:

hash_content.append(f.read())

except OSError:

logger.warning("Failed to read file %s", filepath)

continue

code_hash = hashlib.sha256("\n".join(hash_content).encode()).hexdigest()

# Clear after consumption

self.compilation_config.traced_files.clear()

if not self.compilation_config.cache_dir:

# no provided cache dir, generate one based on the known factors

# that affects the compilation. if none of the factors change,

# the cache dir will be the same so that we can reuse the compiled

# graph.

factors = [env_hash, config_hash, code_hash, compiler_hash]

# Use SHA-256 for cache key hashing to be consistent across

# compute_hash functions. Truncate for a short cache dir name.

hash_key = hashlib.sha256(str(factors).encode()).hexdigest()[:10]

cache_dir = os.path.join(

envs.VLLM_CACHE_ROOT, "torch_compile_cache", hash_key

)

self.compilation_config.cache_dir = cache_dir

We can figure this out as a follow-up (also I'm not sure if you're going to continue working on this or if @zhxchen17 is going to pick up from here)

ProExpertProg

I won't pretend I understand everything but looks overall good to me!

ProExpertProg · 2026-01-15T18:38:56Z

@@ -21,15 +21,17 @@
    "vllm/transformers_utils/config.py",
    "vllm/model_executor/models/registry.py",
    "vllm/compilation/caching.py",
+    "vllm/compilation/piecewise_backend.py",


We should really think about the threat model here, cc @russellb

Signed-off-by: dolpm <34420038+dolpm@users.noreply.github.com>

mergify Bot added ci/build v1 labels Sep 18, 2025

mergify Bot added the needs-rebase label Sep 18, 2025

dolpm force-pushed the remote/serialize-inductor branch from 05baf79 to 96f5d79 Compare September 23, 2025 20:55

mergify Bot added rocm Related to AMD ROCm and removed needs-rebase labels Sep 23, 2025

dolpm force-pushed the remote/serialize-inductor branch 2 times, most recently from 0ce2620 to 38b6517 Compare September 25, 2025 18:31

mergify Bot added the needs-rebase label Sep 25, 2025

dolpm force-pushed the remote/serialize-inductor branch 2 times, most recently from db326a9 to 03fa8b3 Compare September 25, 2025 21:59

mergify Bot removed the needs-rebase label Sep 25, 2025

zhxchen17 reviewed Oct 9, 2025

View reviewed changes

Comment thread vllm/compilation/decorators.py Resolved

Comment thread vllm/compilation/backends.py Resolved

Comment thread vllm/compilation/backends.py Resolved

Comment thread vllm/compilation/backends.py Resolved

Comment thread vllm/compilation/backends.py Resolved

zhxchen17 mentioned this pull request Oct 22, 2025

[CI Failure]: AssertionError on graph pickler while serializing symint nodes. #27348

Closed

1 of 3 tasks

zhxchen17 mentioned this pull request Oct 22, 2025

[compile] Add fallback path to AOT compile when serialization fails. #27350

Merged

3 of 5 tasks

dolpm force-pushed the remote/serialize-inductor branch from 03fa8b3 to 579ec82 Compare October 24, 2025 18:47

mergify Bot added the performance Performance-related issues label Oct 24, 2025

mergify Bot added the needs-rebase label Oct 24, 2025

dolpm force-pushed the remote/serialize-inductor branch from 579ec82 to 869b3fe Compare October 24, 2025 19:26

mergify Bot removed the needs-rebase label Oct 24, 2025

jamesjwu reviewed Oct 24, 2025

View reviewed changes

Comment thread vllm/compilation/compiler_interface.py Resolved

jamesjwu reviewed Oct 24, 2025

View reviewed changes

Comment thread vllm/compilation/caching.py Resolved

jamesjwu reviewed Oct 24, 2025

View reviewed changes

Comment thread vllm/compilation/compiler_interface.py Resolved

zhxchen17 reviewed Oct 24, 2025

View reviewed changes

Comment thread benchmarks/compile/benchmark_inductor_cache.py Resolved

Comment thread vllm/v1/attention/backends/flashinfer.py Resolved

Comment thread vllm/compilation/backends.py Resolved

Comment thread vllm/v1/worker/gpu_worker.py Resolved

dolpm force-pushed the remote/serialize-inductor branch from 869b3fe to 7465ae3 Compare October 27, 2025 17:56