@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 7c5dedc247
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 9, 23:00 UTC

[AOT compilation] support torch.compile inductor artifacts in VllmCompiledFunction (#25205)

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

dolpm committed Jan 20, 2026 at 11:45 UTC 7c5dedc24727904fe144366c2ba68490337460f9
8 files changed +1169 -113
tests/compile/test_aot_compile.py
+415 -2
@@ -2,15 +2,21 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import functools
5 +import hashlib
6 import multiprocessing
7 +import pickle
8 import tempfile
9 from contextlib import contextmanager
10 from pathlib import Path
11 +from unittest.mock import Mock, patch
12
13 import pytest
14 import torch
15
16 import vllm.model_executor.layers.activation
17 +from vllm.compilation.caching import (
18 + StandaloneCompiledArtifacts,
19 +)
20 from vllm.compilation.decorators import support_torch_compile
21 from vllm.config import (
22 CompilationConfig,
@@ -40,6 +46,15 @@ def reference_fn(x: torch.Tensor):
46 return x
47
48
49 +def reference_fn_tuple(x: torch.Tensor):
50 + """Reference function that returns a tuple of tensors."""
51 + assert x.shape[0] <= 42
52 + assert x.shape[0] % 2 == 0
53 + for _ in range(3000):
54 + x = x + x.shape[0]
55 + return x, x * 2
56 +
57 +
58 @support_torch_compile
59 class CompiledMod(torch.nn.Module):
60 def __init__(self, **kwargs):
@@ -49,10 +64,22 @@ class CompiledMod(torch.nn.Module):
64 return reference_fn(x)
65
66
67 +@support_torch_compile
68 +class CompiledModTuple(torch.nn.Module):
69 + """A compiled module that returns a tuple of tensors."""
70 +
71 + def __init__(self, **kwargs):
72 + super().__init__()
73 +
74 + def forward(self, x: torch.Tensor):
75 + return reference_fn_tuple(x)
76 +
77 +
78 def make_vllm_config() -> VllmConfig:
79 return VllmConfig(
80 compilation_config=CompilationConfig(
81 mode=CompilationMode.VLLM_COMPILE,
82 + backend="inductor",
83 )
84 )
85
@@ -73,6 +100,8 @@ def test_no_dynamo_cache_entry(monkeypatch: pytest.MonkeyPatch):
100 expected = reference_fn(*args)
101 with use_vllm_config(vllm_config):
102 m.setenv("VLLM_USE_AOT_COMPILE", "0")
103 + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
104 + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
105 with (
106 pytest.raises(RuntimeError, match="Detected recompile"),
107 torch.compiler.set_stance("fail_on_recompile"),
@@ -94,6 +123,8 @@ def test_force_aot_load(monkeypatch: pytest.MonkeyPatch):
123 with tempfile.TemporaryDirectory() as tmpdirname, monkeypatch.context() as m:
124 args = (torch.randn(10, 10),)
125 m.setenv("VLLM_USE_AOT_COMPILE", "1")
126 + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
127 + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
128 m.setenv("VLLM_FORCE_AOT_LOAD", "1")
129 m.setenv("VLLM_CACHE_ROOT", tmpdirname)
130 vllm_config = make_vllm_config()
@@ -111,18 +142,158 @@ def test_save_and_load(monkeypatch: pytest.MonkeyPatch):
142 with tempfile.TemporaryDirectory() as tmpdirname:
143 m.setenv("VLLM_CACHE_ROOT", tmpdirname)
144 m.setenv("VLLM_USE_AOT_COMPILE", "1")
145 + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
146 + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
147 vllm_config = make_vllm_config()
148 with use_vllm_config(vllm_config):
116 - expected = CompiledMod(vllm_config=vllm_config)(*args)
149 + compiled_mod = CompiledMod(vllm_config=vllm_config)
150 + expected = compiled_mod(*args)
151 +
152 disable_envs_cache()
153
154 m.setenv("VLLM_FORCE_AOT_LOAD", "1")
155 vllm_config = make_vllm_config()
156 with use_vllm_config(vllm_config):
122 - ret = CompiledMod(vllm_config=vllm_config)(*args)
157 + cached_mod = CompiledMod(vllm_config=vllm_config)
158 + ret = cached_mod(*args)
159 + assert cached_mod.was_aot_compile_fn_loaded_from_disk, (
160 + "Expected was_aot_compile_fn_loaded_from_disk to be True"
161 + )
162 assert torch.allclose(ret, expected)
163
164
165 +@pytest.mark.skipif(
166 + not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
167 +)
168 +def test_cache_load_returns_tuple_consistency(monkeypatch: pytest.MonkeyPatch):
169 + """
170 + Test that cache loading correctly handles the returns_tuple logic.
171 +
172 + This verifies that when a model returns a single tensor (not a tuple),
173 + the output type is consistent between fresh compilation and cache load.
174 + Without the fix, cached artifacts would return [tensor] instead of tensor.
175 + """
176 + with monkeypatch.context() as m:
177 + args = (torch.randn(10, 10),)
178 +
179 + with tempfile.TemporaryDirectory() as tmpdirname:
180 + m.setenv("VLLM_CACHE_ROOT", tmpdirname)
181 + m.setenv("VLLM_USE_AOT_COMPILE", "1")
182 + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
183 + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
184 + vllm_config = make_vllm_config()
185 +
186 + # Fresh compilation
187 + with use_vllm_config(vllm_config):
188 + compiled_mod = CompiledMod(vllm_config=vllm_config)
189 + fresh_result = compiled_mod(*args)
190 + fresh_result_type = type(fresh_result)
191 +
192 + # Verify fresh result is a tensor, not a tuple/list
193 + assert isinstance(fresh_result, torch.Tensor), (
194 + f"Fresh compile should return tensor, got {fresh_result_type}"
195 + )
196 +
197 + disable_envs_cache()
198 +
199 + # Load from cache
200 + m.setenv("VLLM_FORCE_AOT_LOAD", "1")
201 + vllm_config = make_vllm_config()
202 + with use_vllm_config(vllm_config):
203 + cached_mod = CompiledMod(vllm_config=vllm_config)
204 + cached_result = cached_mod(*args)
205 + cached_result_type = type(cached_result)
206 +
207 + # Verify cache was actually loaded
208 + assert cached_mod.was_aot_compile_fn_loaded_from_disk, (
209 + "Expected was_aot_compile_fn_loaded_from_disk to be True after "
210 + "loading from cache"
211 + )
212 +
213 + # Verify cached result has same type as fresh result
214 + assert isinstance(cached_result, torch.Tensor), (
215 + f"Cache load should return tensor, got {cached_result_type}. "
216 + "This indicates the returns_tuple logic is not being applied "
217 + "correctly when loading from cache."
218 + )
219 +
220 + # Verify values match
221 + assert torch.allclose(cached_result, fresh_result), (
222 + "Cached result values should match fresh compilation"
223 + )
224 +
225 +
226 +@pytest.mark.skipif(
227 + not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
228 +)
229 +def test_cache_load_returns_tuple_consistency_tuple_output(
230 + monkeypatch: pytest.MonkeyPatch,
231 +):
232 + """
233 + Test that cache loading correctly handles models that return tuples.
234 +
235 + This verifies that when a model returns a tuple of tensors, the output
236 + type is preserved as a tuple between fresh compilation and cache load.
237 + """
238 + with monkeypatch.context() as m:
239 + args = (torch.randn(10, 10),)
240 +
241 + with tempfile.TemporaryDirectory() as tmpdirname:
242 + m.setenv("VLLM_CACHE_ROOT", tmpdirname)
243 + m.setenv("VLLM_USE_AOT_COMPILE", "1")
244 + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
245 + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
246 + vllm_config = make_vllm_config()
247 +
248 + # Fresh compilation with tuple-returning model
249 + with use_vllm_config(vllm_config):
250 + compiled_mod = CompiledModTuple(vllm_config=vllm_config)
251 + fresh_result = compiled_mod(*args)
252 + fresh_result_type = type(fresh_result)
253 +
254 + # Verify fresh result is a tuple
255 + assert isinstance(fresh_result, tuple), (
256 + f"Fresh compile should return tuple, got {fresh_result_type}"
257 + )
258 + assert len(fresh_result) == 2, (
259 + f"Fresh compile should return 2-tuple, got {len(fresh_result)}"
260 + )
261 +
262 + disable_envs_cache()
263 +
264 + # Load from cache
265 + m.setenv("VLLM_FORCE_AOT_LOAD", "1")
266 + vllm_config = make_vllm_config()
267 + with use_vllm_config(vllm_config):
268 + cached_mod = CompiledModTuple(vllm_config=vllm_config)
269 + cached_result = cached_mod(*args)
270 + cached_result_type = type(cached_result)
271 +
272 + # Verify cache was actually loaded
273 + assert cached_mod.was_aot_compile_fn_loaded_from_disk, (
274 + "Expected was_aot_compile_fn_loaded_from_disk to be True after "
275 + "loading from cache"
276 + )
277 +
278 + # Verify cached result is also a tuple
279 + assert isinstance(cached_result, tuple), (
280 + f"Cache load should return tuple, got {cached_result_type}. "
281 + "This indicates the returns_tuple logic is not preserving "
282 + "tuple outputs when loading from cache."
283 + )
284 + assert len(cached_result) == 2, (
285 + f"Cache load should return 2-tuple, got {len(cached_result)}"
286 + )
287 +
288 + # Verify values match
289 + assert torch.allclose(cached_result[0], fresh_result[0]), (
290 + "Cached result[0] values should match fresh compilation"
291 + )
292 + assert torch.allclose(cached_result[1], fresh_result[1]), (
293 + "Cached result[1] values should match fresh compilation"
294 + )
295 +
296 +
297 @pytest.mark.skipif(
298 not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
299 )
@@ -137,6 +308,8 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
308 with tempfile.TemporaryDirectory() as tmpdirname:
309 m.setenv("VLLM_CACHE_ROOT", tmpdirname)
310 m.setenv("VLLM_USE_AOT_COMPILE", "1")
311 + m.setenv("VLLM_USE_MEGA_AOT_ARTIFACT", "1")
312 + m.setenv("VLLM_USE_STANDALONE_COMPILE", "1")
313 vllm_config = make_vllm_config()
314 with use_vllm_config(vllm_config):
315 compiled_mod = CompiledMod(vllm_config=vllm_config)
@@ -144,6 +317,7 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
317 artifacts = compiled_mod.aot_compiled_fn._artifacts
318 guards_string = artifacts.compiled_fn.shape_env.format_guards()
319 assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
320 +
321 disable_envs_cache()
322
323 m.setenv("VLLM_FORCE_AOT_LOAD", "1")
@@ -151,6 +325,9 @@ def test_shape_env(monkeypatch: pytest.MonkeyPatch):
325 with use_vllm_config(vllm_config):
326 compiled_mod = CompiledMod(vllm_config=vllm_config)
327 compiled_mod(*args)
328 + assert compiled_mod.was_aot_compile_fn_loaded_from_disk, (
329 + "Expected was_aot_compile_fn_loaded_from_disk to be True"
330 + )
331 artifacts = compiled_mod.aot_compiled_fn._artifacts
332 guards_string = artifacts.compiled_fn.shape_env.format_guards()
333 assert guards_string == " - s77 <= 42\n - Eq(Mod(s77, 2), 0)"
@@ -188,6 +365,7 @@ def test_partition_wrapper_applied_on_aot_load(
365 with use_vllm_config(vllm_config):
366 compiled_mod = CompiledMod(vllm_config=vllm_config)
367 compiled_mod(*args)
368 +
369 disable_envs_cache()
370
371 # Second run - load from cache, verify partition wrapper applied
@@ -210,6 +388,11 @@ def test_partition_wrapper_applied_on_aot_load(
388 # This tests the fix for the first call after a restart.
389 compiled_mod(*args)
390
391 + # Verify cache was loaded
392 + assert compiled_mod.was_aot_compile_fn_loaded_from_disk, (
393 + "Expected was_aot_compile_fn_loaded_from_disk to be True"
394 + )
395 +
396 # Verify partition wrapper was called on AOT load.
397 assert spy.call_count >= 2, (
398 "Expected partition wrapper to be set and cleared on AOT load, "
@@ -307,3 +490,233 @@ def test_gpt2_cache_hit(monkeypatch: pytest.MonkeyPatch):
490 finally:
491 # Restore original method
492 symbolic_shapes_module.make_symbol = original_make_symbol
493 +
494 +
495 +@pytest.mark.skipif(
496 + not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
497 +)
498 +class TestStandaloneCompiledArtifacts:
499 + def test_init(self):
500 + cache = StandaloneCompiledArtifacts()
501 + assert cache.submodule_bytes == {}
502 + assert cache.submodule_bytes_store == {}
503 + assert cache.loaded_submodule_store == {}
504 +
505 + def test_insert_new_artifact(self):
506 + cache = StandaloneCompiledArtifacts()
507 + test_data = b"test_artifact_data"
508 + submod_name = "test_submod"
509 + shape = "s1"
510 +
511 + hasher = hashlib.sha256()
512 + hasher.update(test_data)
513 + expected_hash = hasher.hexdigest()
514 +
515 + cache.insert(submod_name, shape, test_data)
516 +
517 + assert f"{submod_name}_{shape}" in cache.submodule_bytes
518 + assert cache.submodule_bytes[f"{submod_name}_{shape}"] == expected_hash
519 + assert expected_hash in cache.submodule_bytes_store
520 + assert cache.submodule_bytes_store[expected_hash] == test_data
521 +
522 + def test_insert_duplicate_artifact(self):
523 + cache = StandaloneCompiledArtifacts()
524 +
525 + test_data = b"duplicate_test_data"
526 + submod_name1 = "submod1"
527 + submod_name2 = "submod2"
528 + shape = "s2"
529 +
530 + cache.insert(submod_name1, shape, test_data)
531 + cache.insert(submod_name2, shape, test_data)
532 +
533 + hash1 = cache.submodule_bytes[f"{submod_name1}_{shape}"]
534 + hash2 = cache.submodule_bytes[f"{submod_name2}_{shape}"]
535 + assert hash1 == hash2
536 +
537 + assert len(cache.submodule_bytes_store) == 1
538 + assert len(cache.submodule_bytes) == 2
539 +
540 + def test_get_artifact(self):
541 + cache = StandaloneCompiledArtifacts()
542 + test_data = b"retrievable_data"
543 + submod_name = "mod1"
544 + shape = "shape16"
545 +
546 + cache.insert(submod_name, shape, test_data)
547 + retrieved_data = cache.get(submod_name, shape)
548 +
549 + assert retrieved_data == test_data
550 +
551 + def test_get_nonexistent_artifact(self):
552 + cache = StandaloneCompiledArtifacts()
553 +
554 + with pytest.raises(KeyError):
555 + cache.get("nonexistent", "shape")
556 +
557 + def test_size_bytes(self):
558 + cache = StandaloneCompiledArtifacts()
559 +
560 + assert cache.size_bytes() == 0
561 +
562 + data1 = b"x" * 100
563 + data2 = b"y" * 200
564 + cache.insert("mod1", "shape1", data1)
565 + cache.insert("mod2", "shape2", data2)
566 +
567 + assert cache.size_bytes() == 300
568 +
569 + def test_num_artifacts_and_entries(self):
570 + cache = StandaloneCompiledArtifacts()
571 +
572 + assert cache.num_artifacts() == 0
573 + assert cache.num_entries() == 0
574 +
575 + cache.insert("mod1", "shape1", b"data1")
576 + cache.insert("mod2", "shape2", b"data2")
577 + assert cache.num_artifacts() == 2
578 + assert cache.num_entries() == 2
579 +
580 + cache.insert("mod3", "shape3", b"data1")
581 + assert cache.num_artifacts() == 2
582 + assert cache.num_entries() == 3
583 +
584 + @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize")
585 + def test_load_all_success(self, mock_deserialize):
586 + """Test successful loading of all artifacts"""
587 + cache = StandaloneCompiledArtifacts()
588 +
589 + mock_artifact1 = Mock()
590 + mock_artifact2 = Mock()
591 + mock_deserialize.side_effect = [mock_artifact1, mock_artifact2]
592 +
593 + cache.insert("mod1", "shape1", pickle.dumps(b"data1"))
594 + cache.insert("mod2", "shape2", pickle.dumps(b"data2"))
595 +
596 + cache.load_all()
597 +
598 + assert len(cache.loaded_submodule_store) == 2
599 + assert mock_deserialize.call_count == 2
600 +
601 + @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize")
602 + def test_load_all_already_loaded(self, mock_deserialize):
603 + """Test that load_all skips if already loaded"""
604 + cache = StandaloneCompiledArtifacts()
605 +
606 + mock_artifact = Mock()
607 + cache.submodule_bytes_store["hash1"] = pickle.dumps(b"data1")
608 + cache.loaded_submodule_store["hash1"] = mock_artifact
609 +
610 + cache.load_all()
611 +
612 + mock_deserialize.assert_not_called()
613 +
614 + @patch("torch._inductor.standalone_compile.AOTCompiledArtifact.deserialize")
615 + def test_get_loaded_artifact(self, mock_deserialize):
616 + """Test retrieving loaded artifacts"""
617 + cache = StandaloneCompiledArtifacts()
618 +
619 + mock_artifact = Mock()
620 + mock_deserialize.return_value = mock_artifact
621 +
622 + submod_name = "test_mod"
623 + shape = "test_shape"
624 + cache.insert(submod_name, shape, pickle.dumps(b"test_data"))
625 + cache.load_all()
626 +
627 + retrieved_artifact = cache.get_loaded(submod_name, shape)
628 + assert retrieved_artifact == mock_artifact
629 +
630 + def test_getstate_setstate(self):
631 + cache = StandaloneCompiledArtifacts()
632 +
633 + cache.insert("mod1", "shape1", b"data1")
634 + cache.insert("mod2", "shape2", b"data2")
635 +
636 + cache.loaded_submodule_store["hash1"] = Mock()
637 +
638 + state = cache.__getstate__()
639 +
640 + assert "submodule_bytes" in state
641 + assert "submodule_bytes_store" in state
642 + assert "loaded_submodule_store" not in state
643 +
644 + new_cache = StandaloneCompiledArtifacts()
645 + new_cache.__setstate__(state)
646 +
647 + assert new_cache.submodule_bytes == cache.submodule_bytes
648 + assert new_cache.submodule_bytes_store == cache.submodule_bytes_store
649 + assert new_cache.loaded_submodule_store == {}
650 +
651 + def test_pickle_roundtrip(self):
652 + cache = StandaloneCompiledArtifacts()
653 +
654 + test_data1 = b"pickle_test_data_1"
655 + test_data2 = b"pickle_test_data_2"
656 + cache.insert("mod1", "shape1", test_data1)
657 + cache.insert("mod2", "shape2", test_data2)
658 +
659 + pickled_data = pickle.dumps(cache)
660 + restored_cache = pickle.loads(pickled_data)
661 +
662 + assert restored_cache.get("mod1", "shape1") == test_data1
663 + assert restored_cache.get("mod2", "shape2") == test_data2
664 + assert restored_cache.num_artifacts() == cache.num_artifacts()
665 + assert restored_cache.num_entries() == cache.num_entries()
666 + assert restored_cache.size_bytes() == cache.size_bytes()
667 +
668 + assert len(restored_cache.loaded_submodule_store) == 0
669 +
670 +
671 +@pytest.mark.skipif(
672 + not is_torch_equal_or_newer("2.10.0.dev"), reason="requires torch 2.10"
673 +)
674 +class TestStandaloneCompiledArtifactsIntegration:
675 + def test_add_pickle_unpickle(self):
676 + cache = StandaloneCompiledArtifacts()
677 +
678 + artifacts = {
679 + ("mod1", "shape1"): b"m1s1_artifact",
680 + ("mod1", "shape2"): b"m1s2_artifact",
681 + ("mod2", "shape1"): b"m2s1_artifact",
682 + ("mod2", "shape2"): b"m2s2_artifact",
683 + }
684 +
685 + for (submod, shape), data in artifacts.items():
686 + cache.insert(submod, shape, data)
687 +
688 + assert cache.num_entries() == 4
689 + assert cache.num_artifacts() == 4
690 +
691 + for (submod, shape), expected_data in artifacts.items():
692 + retrieved_data = cache.get(submod, shape)
693 + assert retrieved_data == expected_data
694 +
695 + pickled = pickle.dumps(cache)
696 + restored_cache = pickle.loads(pickled)
697 +
698 + for (submod, shape), expected_data in artifacts.items():
699 + retrieved_data = restored_cache.get(submod, shape)
700 + assert retrieved_data == expected_data
701 +
702 + def test_deduplication(self):
703 + cache = StandaloneCompiledArtifacts()
704 +
705 + shared_data = b"shared_artifact_data" * 1000
706 +
707 + cache.insert("mod1", "shape1", shared_data)
708 + cache.insert("mod2", "shape1", shared_data)
709 + cache.insert("mod1", "shape2", shared_data)
710 + cache.insert("mod3", "shape3", shared_data)
711 +
712 + assert cache.num_entries() == 4
713 + assert cache.num_artifacts() == 1
714 + assert cache.size_bytes() == len(shared_data)
715 +
716 + for submod, shape in [
717 + ("mod1", "shape1"),
718 + ("mod2", "shape1"),
719 + ("mod1", "shape2"),
720 + ("mod3", "shape3"),
721 + ]:
722 + assert cache.get(submod, shape) == shared_data
tools/pre_commit/check_pickle_imports.py
+2
@@ -21,6 +21,7 @@ ALLOWED_FILES = {
21 "vllm/transformers_utils/config.py",
22 "vllm/model_executor/models/registry.py",
23 "vllm/compilation/caching.py",
24 + "vllm/compilation/piecewise_backend.py",
25 "vllm/distributed/utils.py",
26 "vllm/distributed/parallel_state.py",
27 "vllm/distributed/device_communicators/all_reduce_utils.py",
@@ -30,6 +31,7 @@ ALLOWED_FILES = {
31 "tests/multimodal/media/test_base.py",
32 "tests/tokenizers_/test_hf.py",
33 "tests/utils_/test_hashing.py",
34 + "tests/compile/test_aot_compile.py",
35 "benchmarks/kernels/graph_machete_bench.py",
36 "benchmarks/kernels/benchmark_lora.py",
37 "benchmarks/kernels/benchmark_machete.py",
vllm/compilation/backends.py
+228 -60
@@ -2,6 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import ast
5 +import contextvars
6 import dataclasses
7 import hashlib
8 import json
@@ -34,7 +35,6 @@ from vllm.platforms import current_platform
35 from vllm.utils.import_utils import resolve_obj_by_qualname
36 from vllm.utils.torch_utils import is_torch_equal_or_newer
37
37 -from .caching import VllmSerializableFunction
38 from .compiler_interface import (
39 CompilerInterface,
40 EagerAdaptor,
@@ -49,7 +49,48 @@ from .pass_manager import PostGradPassManager
49 logger = init_logger(__name__)
50
51
52 +def make_copy_and_call(
53 + sym_tensor_indices: list[int],
54 + input_buffers: list[torch.Tensor | None],
55 + callable_fn: Callable[..., Any],
56 +) -> Callable[..., Any]:
57 + """Create a wrapper that copies inputs to static buffers before calling.
58 +
59 + This is used for cudagraph input copying where we need to copy dynamic
60 + tensors to static buffers before invoking the compiled graph.
61 +
62 + Args:
63 + sym_tensor_indices: Indices of tensors with symbolic shapes
64 + input_buffers: List of static buffers (can contain None for lazy init)
65 + callable_fn: The compiled function to call
66 +
67 + Returns:
68 + A wrapper function that copies inputs and calls the compiled function
69 + """
70 +
71 + def copy_and_call(*args):
72 + list_args = list(args)
73 + for i, index in enumerate(sym_tensor_indices):
74 + runtime_tensor = list_args[index]
75 + runtime_shape = runtime_tensor.shape[0]
76 +
77 + # lazy initialization of buffer on first call
78 + if input_buffers[i] is None:
79 + input_buffers[i] = runtime_tensor.clone()
80 +
81 + static_tensor = input_buffers[i][:runtime_shape] # type: ignore[index]
82 + static_tensor.copy_(runtime_tensor)
83 + list_args[index] = static_tensor
84 + return callable_fn(*list_args)
85 +
86 + return copy_and_call
87 +
88 +
89 def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
90 + assert not envs.VLLM_USE_MEGA_AOT_ARTIFACT or envs.VLLM_USE_STANDALONE_COMPILE, (
91 + "VLLM_USE_MEGA_AOT_ARTIFACT=1 requires VLLM_USE_STANDALONE_COMPILE=1"
92 + )
93 +
94 if compilation_config.backend == "inductor":
95 # Use standalone compile only if requested, version is new enough,
96 # and the symbol actually exists in this PyTorch build.
@@ -355,6 +396,60 @@ def split_graph(
396 compilation_start_time = 0.0
397
398
399 +def wrap_with_cudagraph_if_needed(
400 + piecewise_backend: Any,
401 + vllm_config: VllmConfig,
402 + compilation_config: CompilationConfig,
403 + is_first_graph: bool,
404 + is_last_graph: bool,
405 +) -> Any:
406 + """
407 + Wrap a piecewise backend with CUDA graph wrapper if needed.
408 + This function is shared between VllmBackend and
409 + construct_serializable_fn_from_inductor_cache.
410 +
411 + Args:
412 + piecewise_backend: The backend to wrap
413 + vllm_config: The vLLM configuration
414 + compilation_config: The compilation configuration
415 + is_first_graph: Whether this is the first graph in the sequence
416 + is_last_graph: Whether this is the last graph in the sequence
417 +
418 + Returns:
419 + The wrapped backend if CUDA graphs are enabled, otherwise the original backend
420 + """
421 + if (
422 + not compilation_config.cudagraph_mode.has_piecewise_cudagraphs()
423 + or compilation_config.use_inductor_graph_partition
424 + ):
425 + return piecewise_backend
426 +
427 + # We're using Dynamo-based piecewise splitting, so we wrap
428 + # the whole subgraph with a static graph wrapper.
429 + from .cuda_graph import CUDAGraphOptions
430 +
431 + # resolve the static graph wrapper class (e.g. CUDAGraphWrapper
432 + # class) as platform dependent.
433 + static_graph_wrapper_class = resolve_obj_by_qualname(
434 + current_platform.get_static_graph_wrapper_cls()
435 + )
436 +
437 + # Always assign PIECEWISE runtime mode to the
438 + # CUDAGraphWrapper for piecewise_backend, to distinguish
439 + # it from the FULL cudagraph runtime mode, no matter it
440 + # is wrapped on a full or piecewise fx graph.
441 + return static_graph_wrapper_class(
442 + runnable=piecewise_backend,
443 + vllm_config=vllm_config,
444 + runtime_mode=CUDAGraphMode.PIECEWISE,
445 + cudagraph_options=CUDAGraphOptions(
446 + debug_log_enable=is_first_graph,
447 + gc_disable=not is_first_graph,
448 + weak_ref_output=is_last_graph,
449 + ),
450 + )
451 +
452 +
453 class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]
454 """Code adapted from `torch.fx.passes.shape_prop.ShapeProp`.
455 It runs the given graph with fake inputs, and compile some
@@ -365,6 +460,18 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]
460 it will be used to determine the order of the compiled piecewise
461 graphs. The first graph will handle logging, and the last graph
462 has some special cudagraph output handling.
463 +
464 + Note: This class shares similar logic with
465 + reconstruct_serializable_fn_from_mega_artifact in caching.py.
466 + Both create PiecewiseBackend instances and wrap them with cudagraph.
467 + The key difference is:
468 + - reconstruct_serializable_fn_from_mega_artifact: PiecewiseBackend receives
469 + pre-compiled runnables (compiled_runnables is set, graph is None)
470 + - this class: PiecewiseBackend receives the FX graph to compile
471 + (graph is set, compiled_runnables is None)
472 +
473 +
474 + If modifying the backend creation/wrapping logic, consider updating both.
475 """
476
477 def __init__(
@@ -413,6 +520,8 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]
520 ]
521
522 # Lazy import here to avoid circular import
523 + from torch._inductor.compile_fx import graph_returns_tuple
524 +
525 from .piecewise_backend import PiecewiseBackend
526
527 piecewise_backend = PiecewiseBackend(
@@ -422,38 +531,16 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]
531 len(self.compile_submod_names),
532 sym_shape_indices,
533 self.vllm_backend,
534 + graph_returns_tuple(submod),
535 )
536
427 - if (
428 - self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs()
429 - and not self.compilation_config.use_inductor_graph_partition
430 - ):
431 - # We're using Dynamo-based piecewise splitting, so we wrap
432 - # the whole subgraph with a static graph wrapper.
433 - from .cuda_graph import CUDAGraphOptions
434 -
435 - # resolve the static graph wrapper class (e.g. CUDAGraphWrapper
436 - # class) as platform dependent.
437 - static_graph_wrapper_class = resolve_obj_by_qualname(
438 - current_platform.get_static_graph_wrapper_cls()
439 - )
440 -
441 - # Always assign PIECEWISE runtime mode to the
442 - # CUDAGraphWrapper for piecewise_backend, to distinguish
443 - # it from the FULL cudagraph runtime mode, no matter it
444 - # is wrapped on a full or piecewise fx graph.
445 - self.module.__dict__[target] = static_graph_wrapper_class(
446 - runnable=piecewise_backend,
447 - vllm_config=self.vllm_config,
448 - runtime_mode=CUDAGraphMode.PIECEWISE,
449 - cudagraph_options=CUDAGraphOptions(
450 - debug_log_enable=piecewise_backend.is_first_graph,
451 - gc_disable=not piecewise_backend.is_first_graph,
452 - weak_ref_output=piecewise_backend.is_last_graph,
453 - ),
454 - )
455 - else:
456 - self.module.__dict__[target] = piecewise_backend
537 + self.module.__dict__[target] = wrap_with_cudagraph_if_needed(
538 + piecewise_backend,
539 + self.vllm_config,
540 + self.compilation_config,
541 + piecewise_backend.is_first_graph,
542 + piecewise_backend.is_last_graph,
543 + )
544
545 compilation_counter.num_piecewise_capturable_graphs_seen += 1
546
@@ -465,6 +552,21 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]
552 model_tag: str = "backbone"
553 model_is_encoder: bool = False
554
555 +_on_compilation_complete_callback: contextvars.ContextVar[Callable[[], None] | None] = (
556 + contextvars.ContextVar("on_compilation_complete_callback", default=None)
557 +)
558 +
559 +
560 +@contextmanager
561 +def set_on_compilation_complete(
562 + callback: Callable[[], None],
563 +) -> Generator[None, None, None]:
564 + token = _on_compilation_complete_callback.set(callback)
565 + try:
566 + yield
567 + finally:
568 + _on_compilation_complete_callback.reset(token)
569 +
570
571 @contextmanager
572 def set_model_tag(tag: str, is_encoder: bool = False) -> Generator[None, None, None]:
@@ -509,8 +611,6 @@ class VllmBackend:
611 returned_callable: Callable[..., Any]
612 # Inductor passes to run on the graph pre-defunctionalization
613 post_grad_passes: Sequence[Callable[..., Any]]
512 - sym_tensor_indices: list[int]
513 - input_buffers: list[torch.Tensor]
614 compiler_manager: CompilerManager
615 # Copy of CompilationConfig.inductor_compile_config +
616 # an entry for PostGradPassManager
@@ -539,9 +639,6 @@ class VllmBackend:
639 )()
640 self.pass_key = current_platform.pass_key
641
542 - self.sym_tensor_indices = []
543 - self.input_buffers = []
544 -
642 self.vllm_config = vllm_config
643 self.compilation_config = vllm_config.compilation_config
644
@@ -558,6 +655,68 @@ class VllmBackend:
655 # `torch.compile` is JIT compiled, so we don't need to
656 # do anything here
657
658 + def collect_standalone_compile_artifacts(
659 + self,
660 + ) -> tuple[Any, dict[str, list[int]] | None, dict[str, bool] | None]:
661 + """Collect inductor cache artifacts from all piecewise backends.
662 +
663 + Returns:
664 + tuple: (standalone_compile_artifacts, sym_shape_indices_map,
665 + returns_tuple_map)
666 + - standalone_compile_artifacts: StandaloneCompiledArtifacts
667 + with compiled artifacts
668 + - sym_shape_indices_map: dict mapping submod_name to
669 + sym_shape_indices
670 + - returns_tuple_map: dict mapping submod_name to
671 + returns_tuple
672 + """
673 +
674 + if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
675 + return None, None, None
676 +
677 + from .caching import StandaloneCompiledArtifacts
678 + from .piecewise_backend import PiecewiseBackend
679 +
680 + standalone_compile_artifacts = StandaloneCompiledArtifacts()
681 + sym_shape_indices_map = {}
682 + returns_tuple_map = {}
683 +
684 + for name, _ in self.split_gm.named_children():
685 + # get the actual attribute (shadowed by PiecewiseBackend in __dict__)
686 + child = getattr(self.split_gm, name)
687 + # unwrap the static graph wrapper class if applicable
688 + piecewise_backend = child.runnable if hasattr(child, "runnable") else child
689 +
690 + if not isinstance(piecewise_backend, PiecewiseBackend):
691 + continue
692 +
693 + submod_name = name
694 + sym_shape_indices_map[submod_name] = piecewise_backend.sym_shape_indices
695 + returns_tuple_map[submod_name] = piecewise_backend.returns_tuple
696 +
697 + for shape_str, bytes_data in piecewise_backend.to_bytes().items():
698 + standalone_compile_artifacts.insert(submod_name, shape_str, bytes_data)
699 + logger.debug(
700 + "collected artifact for %s shape %s (%d bytes)",
701 + submod_name,
702 + shape_str,
703 + len(bytes_data),
704 + )
705 +
706 + logger.info(
707 + "collected artifacts: %d entries, %d artifacts, %d bytes total",
708 + standalone_compile_artifacts.num_entries(),
709 + standalone_compile_artifacts.num_artifacts(),
710 + standalone_compile_artifacts.size_bytes(),
711 + )
712 +
713 + logger.debug(
714 + "standalone compile artifact keys: %s",
715 + list(standalone_compile_artifacts.submodule_bytes.keys()),
716 + )
717 +
718 + return standalone_compile_artifacts, sym_shape_indices_map, returns_tuple_map
719 +
720 def configure_post_pass(self) -> None:
721 self.pass_manager.configure(self.vllm_config)
722
@@ -579,9 +738,11 @@ class VllmBackend:
738 )
739 self.inductor_config[self.pass_key] = self.pass_manager
740
582 - def __call__(
583 - self, graph: fx.GraphModule, example_inputs: Sequence[Any]
584 - ) -> VllmSerializableFunction:
741 + def __call__(self, graph: fx.GraphModule, example_inputs: Sequence[Any]) -> Any:
742 + from .caching import (
743 + VllmSerializableFunction,
744 + )
745 +
746 vllm_config = self.vllm_config
747 # Minimal hashing here with existing utilities, reused below.
748
@@ -721,6 +882,12 @@ class VllmBackend:
882
883 self.split_gm, self.piecewise_graphs = split_graph(graph, fx_split_ops)
884
885 + # keep a split_gm copy from BEFORE the interpreter replaces
886 + # submodules with PiecewiseBackend -- used for serialization
887 + original_split_gm = None
888 + if envs.VLLM_USE_MEGA_AOT_ARTIFACT:
889 + original_split_gm = deepcopy(self.split_gm)
890 +
891 from torch._dynamo.utils import lazy_format_graph_code
892
893 # depyf will hook lazy_format_graph_code and dump the graph
@@ -792,13 +959,21 @@ class VllmBackend:
959 )
960
961 self._called = True
962 + graph_to_serialize = (
963 + original_split_gm if envs.VLLM_USE_MEGA_AOT_ARTIFACT else self.graph
964 + )
965
966 if (
967 self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
968 or not self.compilation_config.cudagraph_copy_inputs
969 ):
970 return VllmSerializableFunction(
801 - graph, example_inputs, self.prefix, self.split_gm, self.is_encoder
971 + graph_to_serialize,
972 + example_inputs,
973 + self.prefix,
974 + self.split_gm,
975 + is_encoder=self.is_encoder,
976 + vllm_backend=self,
977 )
978
979 # index of tensors that have symbolic shapes (batch size)
@@ -806,7 +981,7 @@ class VllmBackend:
981 # symbolic shape only happens for input tensors.
982 from torch.fx.experimental.symbolic_shapes import is_symbolic
983
809 - self.sym_tensor_indices = [
984 + sym_tensor_indices = [
985 i
986 for i, x in enumerate(fake_args)
987 if isinstance(x, torch._subclasses.fake_tensor.FakeTensor)
@@ -816,25 +991,18 @@ class VllmBackend:
991 # compiler managed cudagraph input buffers
992 # we assume the first run with symbolic shapes
993 # has the maximum size among all the tensors
819 - self.input_buffers = [
820 - example_inputs[x].clone() for x in self.sym_tensor_indices
821 - ]
822 -
823 - # this is the callable we return to Dynamo to run
824 - def copy_and_call(*args: Any) -> Any:
825 - list_args = list(args)
826 - for i, index in enumerate(self.sym_tensor_indices):
827 - runtime_tensor = list_args[index]
828 - runtime_shape = runtime_tensor.shape[0]
829 - static_tensor = self.input_buffers[i][:runtime_shape]
830 -
831 - # copy the tensor to the static buffer
832 - static_tensor.copy_(runtime_tensor)
833 -
834 - # replace the tensor in the list_args to the static buffer
835 - list_args[index] = static_tensor
836 - return self.split_gm(*list_args)
994 + copy_and_call = make_copy_and_call(
995 + sym_tensor_indices,
996 + [example_inputs[x].clone() for x in sym_tensor_indices],
997 + self.split_gm,
998 + )
999
1000 return VllmSerializableFunction(
839 - graph, example_inputs, self.prefix, copy_and_call, self.is_encoder
1001 + graph_to_serialize,
1002 + example_inputs,
1003 + self.prefix,
1004 + copy_and_call,
1005 + is_encoder=self.is_encoder,
1006 + vllm_backend=self,
1007 + sym_tensor_indices=sym_tensor_indices,
1008 )
vllm/compilation/caching.py
+327 -8
@@ -1,6 +1,7 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 +import hashlib
5 import inspect
6 import os
7 import pickle
@@ -12,6 +13,7 @@ import torch
13 from torch.utils import _pytree as pytree
14
15 import vllm.envs as envs
16 +from vllm.compilation.compiler_interface import get_inductor_factors
17 from vllm.config import VllmConfig, get_current_vllm_config
18 from vllm.config.utils import hash_factors
19 from vllm.logger import init_logger
@@ -27,6 +29,121 @@ assert isinstance(SerializableCallable, type)
29 logger = init_logger(__name__)
30
31
32 +class StandaloneCompiledArtifacts:
33 + """Storage for standalone compiled artifacts with content-based deduplication.
34 +
35 + Deduplication works via a two-level indirection:
36 + 1. `submodule_bytes` maps "{submod_name}_{shape}" -> SHA256 hash
37 + 2. `submodule_bytes_store` maps SHA256 hash -> actual bytes
38 +
39 + When inserting, we compute the SHA256 hash of the bytes. If the hash
40 + already exists in `submodule_bytes_store`, we reuse the existing entry
41 + rather than storing duplicate bytes. This is common because submodules
42 + often compile to identical artifacts (e.g., identical transformer layers
43 + split on attn)
44 + """
45 +
46 + def __init__(self):
47 + # dict from submodule name to byte hash
48 + self.submodule_bytes = {}
49 + # dict from byte hash to bytes
50 + self.submodule_bytes_store = {}
51 + # dict from byte hash to loaded module
52 + self.loaded_submodule_store = {}
53 +
54 + def insert(self, submod_name: str, shape: str, entry: bytes):
55 + hasher = hashlib.sha256()
56 + hasher.update(entry)
57 + hex_digest = hasher.hexdigest()
58 + self.submodule_bytes[f"{submod_name}_{shape}"] = hex_digest
59 + if hex_digest not in self.submodule_bytes_store:
60 + self.submodule_bytes_store[hex_digest] = entry
61 + logger.debug(
62 + "inserting new artifact for submod %s with shape %s "
63 + "(%s bytes) at hash %s",
64 + submod_name,
65 + shape,
66 + len(entry),
67 + hex_digest,
68 + )
69 + else:
70 + logger.debug(
71 + "reusing existing cache artifact for submod %s "
72 + "with shape %s (%s bytes) at hash %s",
73 + submod_name,
74 + shape,
75 + len(entry),
76 + hex_digest,
77 + )
78 +
79 + def get(self, submod_name: str, shape: str) -> bytes:
80 + logger.debug(
81 + "getting artifact for submod %s with shape %s",
82 + submod_name,
83 + shape,
84 + )
85 + return self.submodule_bytes_store[
86 + self.submodule_bytes[f"{submod_name}_{shape}"]
87 + ]
88 +
89 + def get_loaded(self, submod_name: str, shape: str):
90 + logger.debug(
91 + "getting artifact for submod %s with shape %s",
92 + submod_name,
93 + shape,
94 + )
95 + return self.loaded_submodule_store[
96 + self.submodule_bytes[f"{submod_name}_{shape}"]
97 + ]
98 +
99 + def size_bytes(self) -> int:
100 + return sum(len(entry) for entry in self.submodule_bytes_store.values())
101 +
102 + def num_artifacts(self) -> int:
103 + return len(self.submodule_bytes_store)
104 +
105 + def num_entries(self) -> int:
106 + return len(self.submodule_bytes)
107 +
108 + def submodule_names(self) -> list[str]:
109 + # get unique "{submod_name}" from "{submod_name}_{shape}", preserving order
110 + names = [cache_key.rsplit("_", 1)[0] for cache_key in self.submodule_bytes]
111 + return list(dict.fromkeys(names))
112 +
113 + def load_all(self) -> None:
114 + import concurrent.futures
115 +
116 + # check already loaded
117 + if len(self.loaded_submodule_store) == len(self.submodule_bytes_store):
118 + return
119 +
120 + from torch._inductor.standalone_compile import AOTCompiledArtifact
121 +
122 + def _load_entry(entry_bytes) -> AOTCompiledArtifact:
123 + entry = pickle.loads(entry_bytes)
124 + return AOTCompiledArtifact.deserialize(entry)
125 +
126 + with concurrent.futures.ThreadPoolExecutor() as executor:
127 + entries = list(self.submodule_bytes_store.values())
128 + loaded_entries = list(executor.map(_load_entry, entries))
129 +
130 + for i, k in enumerate(self.submodule_bytes_store.keys()):
131 + self.loaded_submodule_store[k] = loaded_entries[i]
132 +
133 + logger.debug("loaded all %s submodules", self.num_artifacts())
134 +
135 + def __getstate__(self):
136 + return {
137 + "submodule_bytes": self.submodule_bytes,
138 + "submodule_bytes_store": self.submodule_bytes_store,
139 + }
140 +
141 + def __setstate__(self, state):
142 + self.submodule_bytes = state["submodule_bytes"]
143 + self.submodule_bytes_store = state["submodule_bytes_store"]
144 + self.loaded_submodule_store = {}
145 +
146 +
147 class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
148 """
149 A wrapper around a compiled function by vllm. It will forward the tensor
@@ -46,6 +163,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
163 prefix: str,
164 optimized_call: Callable[..., Any],
165 is_encoder: bool = False,
166 + vllm_backend: Any | None = None,
167 + sym_tensor_indices: list[int] | None = None,
168 ) -> None:
169 assert isinstance(graph_module, torch.fx.GraphModule)
170 self.graph_module = graph_module
@@ -54,6 +173,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
173 self.optimized_call = optimized_call
174 self.is_encoder = is_encoder
175 self.shape_env = None
176 + self.vllm_backend = vllm_backend
177 + self.sym_tensor_indices = sym_tensor_indices
178 sym_input = next(
179 (i for i in self.example_inputs if isinstance(i, torch.SymInt)), None
180 )
@@ -74,9 +195,15 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
195 state = compiled_fn.__dict__.copy()
196 state.pop("optimized_call")
197 state.pop("shape_env")
198 + state.pop("vllm_backend", None)
199 for node in state["graph_module"].graph.nodes:
200 node.meta.pop("source_fn_stack", None)
201 node.meta.pop("nn_module_stack", None)
202 + for name, submod in state["graph_module"].named_children():
203 + if hasattr(submod, "graph"):
204 + for node in submod.graph.nodes:
205 + node.meta.pop("source_fn_stack", None)
206 + node.meta.pop("nn_module_stack", None)
207
208 graph_reducer_override = GraphPickler.reducer_override
209
@@ -93,15 +220,36 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
220 return type(None), ()
221 return graph_reducer_override(self, obj)
222
96 - # Mask off tensor inputs since they are large and not needed.
97 - state["example_inputs"] = pytree.tree_map_only(
98 - torch.Tensor, lambda _: None, state["example_inputs"]
99 - )
223 + if state.get("sym_tensor_indices"):
224 + # put tensor inputs on meta device since their data
225 + # isn't needed, yet we need the meta for make_copy_and_call
226 + state["example_inputs"] = pytree.tree_map_only(
227 + torch.Tensor,
228 + lambda inp: torch.empty_like(inp, device="meta"),
229 + state["example_inputs"],
230 + )
231 + else:
232 + # mask off all tensor inputs since they are large and not needed.
233 + state["example_inputs"] = pytree.tree_map_only(
234 + torch.Tensor,
235 + lambda inp: torch.empty_like(inp, device="meta"),
236 + state["example_inputs"],
237 + )
238 with patch.object(GraphPickler, "reducer_override", _graph_reducer_override):
239 state["graph_module"] = GraphPickler.dumps(
240 state["graph_module"], Options(ops_filter=None)
241 )
242 state["example_inputs"] = GraphPickler.dumps(state["example_inputs"])
243 +
244 + if compiled_fn.vllm_backend:
245 + (
246 + standalone_compile_artifacts,
247 + sym_shape_indices_map,
248 + returns_tuple_map,
249 + ) = compiled_fn.vllm_backend.collect_standalone_compile_artifacts()
250 + state["standalone_compile_artifacts"] = standalone_compile_artifacts
251 + state["sym_shape_indices_map"] = sym_shape_indices_map
252 + state["returns_tuple_map"] = returns_tuple_map
253 return pickle.dumps(state)
254
255 @classmethod
@@ -111,15 +259,48 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
259 from torch.fx._graph_pickler import GraphPickler
260 from torch.fx.experimental.symbolic_shapes import ShapeEnv
261
114 - from vllm.compilation.backends import VllmBackend
115 -
262 state = pickle.loads(data)
263 fake_mode = FakeTensorMode(shape_env=ShapeEnv())
264 state["graph_module"] = GraphPickler.loads(state["graph_module"], fake_mode)
265 state["graph_module"].recompile()
266 state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode)
267 +
268 + standalone_compile_artifacts = state.pop("standalone_compile_artifacts", None)
269 + sym_shape_indices_map = state.pop("sym_shape_indices_map", {})
270 + returns_tuple_map = state.pop("returns_tuple_map", {})
271 +
272 + if envs.VLLM_USE_MEGA_AOT_ARTIFACT:
273 + assert standalone_compile_artifacts is not None
274 + submod_names = standalone_compile_artifacts.submodule_names()
275 + num_submods = len(submod_names)
276 + num_artifacts = standalone_compile_artifacts.num_artifacts()
277 +
278 + logger.info(
279 + "reconstructing serializable fn from standalone compile "
280 + "artifacts. num_artifacts=%d num_submods=%d",
281 + num_artifacts,
282 + num_submods,
283 + )
284 +
285 + fn = reconstruct_serializable_fn_from_mega_artifact(
286 + state=state,
287 + standalone_compile_artifacts=standalone_compile_artifacts,
288 + vllm_config=get_current_vllm_config(),
289 + sym_shape_indices_map=sym_shape_indices_map,
290 + returns_tuple_map=returns_tuple_map,
291 + )
292 +
293 + logger.info(
294 + "reconstructed serializable fn from standalone compile artifacts"
295 + )
296 +
297 + return fn
298 +
299 + # Fall back to standard VllmBackend
300 + from vllm.compilation.backends import VllmBackend
301 +
302 is_encoder = state.get("is_encoder", False)
122 - vllm_backend = VllmBackend(
303 + vllm_backend: VllmBackend = VllmBackend(
304 get_current_vllm_config(), state["prefix"], is_encoder
305 )
306
@@ -152,7 +333,140 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
333 return "VllmSerializableFunction"
334
335
155 -def compilation_config_hash_factors(vllm_config: VllmConfig) -> list[str]:
336 +def reconstruct_serializable_fn_from_mega_artifact(
337 + state: dict[str, Any],
338 + standalone_compile_artifacts: "StandaloneCompiledArtifacts",
339 + vllm_config: VllmConfig,
340 + sym_shape_indices_map: dict[str, list[int]],
341 + returns_tuple_map: dict[str, bool],
342 +) -> "VllmSerializableFunction":
343 + """Construct a VllmSerializableFunction from cached inductor artifacts.
344 +
345 + This function reconstructs a callable model from pre-compiled inductor
346 + artifacts without re-running the compilation. It:
347 + 1. Loads all cached artifacts
348 + 2. Builds compiled callables for each submodule/shape
349 + 3. Creates PiecewiseBackend instances that dispatch to cached artifacts
350 + 4. Wraps with cudagraph if needed
351 + 5. Returns the final VllmSerializableFunction
352 +
353 + Note: This function shares similar logic with PiecewiseCompileInterpreter
354 + in backends.py. Both create PiecewiseBackend instances and wrap them with
355 + cudagraph. The key difference is:
356 + - this function: PiecewiseBackend receives pre-compiled runnables
357 + (compiled_runnables is set, graph is None)
358 + - PiecewiseCompileInterpreter: PiecewiseBackend receives the FX graph
359 + to compile (graph is set, compiled_runnables is None)
360 +
361 + If modifying the backend creation/wrapping logic, consider updating both.
362 +
363 + Args:
364 + state: Deserialized state dict containing graph_module, example_inputs,
365 + prefix, sym_tensor_indices, is_encoder, etc.
366 + standalone_compile_artifacts: The StandaloneCompiledArtifacts containing
367 + pre-compiled artifacts for each submodule/shape combination.
368 + vllm_config: The vLLM configuration.
369 + sym_shape_indices_map: Mapping from submod_name to sym_shape_indices.
370 + returns_tuple_map: Mapping from submod_name to returns_tuple.
371 +
372 + Returns:
373 + A VllmSerializableFunction that can be called directly.
374 + """
375 + from vllm.compilation.backends import (
376 + VllmBackend,
377 + make_copy_and_call,
378 + wrap_with_cudagraph_if_needed,
379 + )
380 + from vllm.compilation.piecewise_backend import PiecewiseBackend
381 +
382 + prefix = state["prefix"]
383 + is_encoder = state.get("is_encoder", False)
384 + split_gm = state["graph_module"]
385 + compilation_config = vllm_config.compilation_config
386 +
387 + standalone_compile_artifacts.load_all()
388 +
389 + submod_names = standalone_compile_artifacts.submodule_names()
390 + compiled_callables: dict[str, dict[str, Callable]] = {}
391 +
392 + for cache_key in standalone_compile_artifacts.submodule_bytes:
393 + submod_name, shape_str = cache_key.rsplit("_", 1)
394 + compiled_callables.setdefault(submod_name, {})[shape_str] = (
395 + standalone_compile_artifacts.get_loaded(submod_name, shape_str)
396 + )
397 +
398 + vllm_backend = VllmBackend(vllm_config, prefix, is_encoder)
399 + dummy_cache_dir = os.path.join(envs.VLLM_CACHE_ROOT, "dummy_cache")
400 + os.makedirs(dummy_cache_dir, exist_ok=True)
401 + vllm_backend.compiler_manager.initialize_cache(
402 + cache_dir=dummy_cache_dir,
403 + disable_cache=True,
404 + prefix=prefix,
405 + )
406 +
407 + # spot check that cached submodules exist in the graph structure
408 + graph_children = {name for name, _ in split_gm.named_children()}
409 + missing = set(submod_names) - graph_children
410 + assert not missing, (
411 + f"artifacts reference submodules not in graph: {missing}. "
412 + f"graph has: {sorted(graph_children)}"
413 + )
414 +
415 + for i, submod_name in enumerate(submod_names):
416 + assert submod_name in sym_shape_indices_map and submod_name in returns_tuple_map
417 +
418 + sym_shape_indices = sym_shape_indices_map[submod_name]
419 + returns_tuple = returns_tuple_map[submod_name]
420 + runnables = compiled_callables[submod_name]
421 +
422 + piecewise_backend = PiecewiseBackend(
423 + graph=None, # not needed for cached artifacts
424 + vllm_config=vllm_config,
425 + piecewise_compile_index=i,
426 + total_piecewise_compiles=len(submod_names),
427 + sym_shape_indices=sym_shape_indices,
428 + vllm_backend=vllm_backend,
429 + returns_tuple=returns_tuple,
430 + compiled_runnables=runnables,
431 + )
432 +
433 + is_first = i == 0
434 + is_last = i == len(submod_names) - 1
435 + wrapped_backend = wrap_with_cudagraph_if_needed(
436 + piecewise_backend,
437 + vllm_config,
438 + compilation_config,
439 + is_first,
440 + is_last,
441 + )
442 +
443 + split_gm.__dict__[submod_name] = wrapped_backend
444 + logger.debug(
445 + "Replaced submodule %s with piecewise backend from cache",
446 + submod_name,
447 + )
448 +
449 + if compilation_config.cudagraph_copy_inputs:
450 + sym_tensor_indices = state["sym_tensor_indices"]
451 + input_buffers = [
452 + torch.empty_like(
453 + state["example_inputs"][idx], device=vllm_config.device_config.device
454 + )
455 + for idx in sym_tensor_indices
456 + ]
457 + optimized_call = make_copy_and_call(sym_tensor_indices, input_buffers, split_gm)
458 + else:
459 + optimized_call = split_gm
460 +
461 + fn = VllmSerializableFunction(
462 + **state,
463 + optimized_call=optimized_call,
464 + vllm_backend=None,
465 + )
466 + return fn
467 +
468 +
469 +def aot_compile_hash_factors(vllm_config: VllmConfig) -> list[str]:
470 factors = []
471 # 0. factors come from the env, for example, The values of
472 # VLLM_PP_LAYER_PARTITION will affect the computation graph.
@@ -163,6 +477,11 @@ def compilation_config_hash_factors(vllm_config: VllmConfig) -> list[str]:
477 # model is created)
478 config_hash = vllm_config.compute_hash()
479 factors.append(config_hash)
480 +
481 + # 2. inductor factors if applicable
482 + if envs.VLLM_USE_MEGA_AOT_ARTIFACT:
483 + factors.extend(get_inductor_factors())
484 +
485 return factors
486
487
vllm/compilation/compiler_interface.py
+41 -7
@@ -16,9 +16,12 @@ import vllm.envs as envs
16 from vllm.compilation.counter import compilation_counter
17 from vllm.config import VllmConfig
18 from vllm.config.utils import Range
19 +from vllm.logger import init_logger
20 from vllm.utils.hashing import safe_hash
21 from vllm.utils.torch_utils import is_torch_equal_or_newer
22
23 +logger = init_logger(__name__)
24 +
25
26 class CompilerInterface:
27 """
@@ -230,12 +233,42 @@ class InductorStandaloneAdaptor(CompilerInterface):
233
234 from torch._inductor import standalone_compile
235
233 - compiled_graph = standalone_compile(
234 - graph,
235 - example_inputs,
236 - dynamic_shapes=dynamic_shapes,
237 - options={"config_patches": current_config},
238 - )
236 + supports_aot = is_torch_equal_or_newer("2.10.0.dev")
237 +
238 + if not supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT:
239 + logger.error(
240 + "CRITICAL: VLLM_USE_MEGA_AOT_ARTIFACT "
241 + "is enabled but PyTorch version does not support 'aot' "
242 + "parameter in standalone_compile. This requires PyTorch "
243 + "2.10.0+. Falling back to non-AOT mode."
244 + )
245 +
246 + compile_kwargs = {
247 + "dynamic_shapes": dynamic_shapes,
248 + "options": {
249 + "config_patches": current_config,
250 + },
251 + }
252 +
253 + use_aot: bool = supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT
254 + # only add 'aot' parameter if both supported and enabled...
255 + # this will set bundled_autograd_cache
256 + # https://github.com/pytorch/pytorch/blob/9bbc5b2905c260adf41bc866a732f9c121a2828a/torch/_inductor/standalone_compile.py#L359 # noqa
257 + if use_aot:
258 + compile_kwargs["aot"] = True # type: ignore[assignment]
259 +
260 + compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
261 +
262 + if use_aot:
263 + from torch._inductor.standalone_compile import AOTCompiledArtifact
264 +
265 + assert isinstance(compiled_graph, AOTCompiledArtifact)
266 + assert hasattr(compiled_graph, "serialize")
267 + # just return the compiled graph and a key
268 + # since we can serialize the bytes using to_bytes
269 + # and reload it using the key when reading
270 + return compiled_graph, None
271 +
272 # Save the compiled artifact to disk in the specified path
273 assert key is not None
274 path = os.path.join(self.cache_dir, key)
@@ -619,7 +652,8 @@ def set_inductor_config(config: dict[str, Any], compile_range: Range) -> None:
652
653
654 def set_functorch_config() -> None:
622 - torch._functorch.config.bundled_autograd_cache = False
655 + if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
656 + torch._functorch.config.bundled_autograd_cache = False
657
658
659 class EagerAdaptor(CompilerInterface):
vllm/compilation/decorators.py
+36 -16
@@ -320,7 +320,7 @@ def _support_torch_compile(
320 return
321
322 self._check_shape_invariants = shape_invariants
323 -
323 + self.was_aot_compile_fn_loaded_from_disk = False
324 compilation_counter.num_models_seen += 1
325 self.compiled = False
326
@@ -417,9 +417,9 @@ def _support_torch_compile(
417 serialized backend artifacts), then we need to generate a new AOT
418 compile artifact from scratch.
419 """
420 - from .caching import compilation_config_hash_factors
420 + from .caching import aot_compile_hash_factors
421
422 - factors: list[str] = compilation_config_hash_factors(self.vllm_config)
422 + factors: list[str] = aot_compile_hash_factors(self.vllm_config)
423
424 factors.append(_model_hash_key(self.forward))
425 hash_key = hashlib.sha256(str(factors).encode()).hexdigest()
@@ -446,6 +446,7 @@ def _support_torch_compile(
446 if not self.compilation_config.dynamic_shapes_config.evaluate_guards:
447 loaded_fn.disable_guard_check()
448 self.aot_compiled_fn = loaded_fn
449 + self.was_aot_compile_fn_loaded_from_disk = True
450 except Exception as e:
451 if os.path.exists(aot_compilation_path):
452 logger.warning(
@@ -547,26 +548,45 @@ def _support_torch_compile(
548 logger.warning("Detected eager backend, disabling AOT compile.")
549 use_aot_compile = False
550 if use_aot_compile:
550 - self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
551 - output = self.aot_compiled_fn(self, *args, **kwargs)
552 - assert aot_compilation_path is not None
553 - assert cache_dir is not None
554 - try:
555 - os.makedirs(cache_dir, exist_ok=True)
556 - self.aot_compiled_fn.save_compiled_function(aot_compilation_path)
557 - except Exception as e:
558 - logger.warning(
559 - "Cannot save aot compilation to path %s, error: %s",
560 - aot_compilation_path,
561 - str(e),
562 - )
551 + from vllm.compilation.backends import set_on_compilation_complete
552 +
553 + # store the path for saving after warmup
554 + self._aot_compilation_path = aot_compilation_path
555 + self._aot_cache_dir = cache_dir
556 + # set callback in context so it's available when compilation completes
557 + with set_on_compilation_complete(self.save_aot_compiled_function):
558 + self.aot_compiled_fn = self.aot_compile(*args, **kwargs)
559 + output = self.aot_compiled_fn(self, *args, **kwargs)
560 else:
561 output = TorchCompileWithNoGuardsWrapper.__call__(self, *args, **kwargs) # type: ignore[arg-type]
562
563 self.compiled = True
564 return output
565
566 + # triggers VllmSerializableFunction.serialize()
567 + def save_aot_compiled_function(self):
568 + if self.was_aot_compile_fn_loaded_from_disk:
569 + logger.debug("AOT compiled function was loaded from cache, skipping save")
570 + return
571 +
572 + assert (
573 + self.aot_compiled_fn and self._aot_compilation_path and self._aot_cache_dir
574 + )
575 +
576 + logger.info("saving AOT compiled function to %s", self._aot_compilation_path)
577 + try:
578 + os.makedirs(self._aot_cache_dir, exist_ok=True)
579 + self.aot_compiled_fn.save_compiled_function(self._aot_compilation_path)
580 + logger.info("saved AOT compiled function to %s", self._aot_compilation_path)
581 + except Exception as e:
582 + logger.warning(
583 + "unable to save AOT compiled function to %s: %s",
584 + self._aot_compilation_path,
585 + e,
586 + )
587 +
588 cls.__call__ = __call__
589 + cls.save_aot_compiled_function = save_aot_compiled_function
590 return cls
591
592
vllm/compilation/piecewise_backend.py
+112 -20
@@ -2,10 +2,15 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 import dataclasses
5 +import io
6 +import pickle
7 from collections.abc import Callable
8 +from pickle import Pickler
9 from typing import Any
10
11 +import torch._functorch.config
12 import torch.fx as fx
13 +from torch._inductor.runtime.triton_heuristics import CachingAutotuner
14
15 from vllm.compilation.backends import VllmBackend
16 from vllm.compilation.monitor import end_monitoring_torch_compile
@@ -26,12 +31,14 @@ class RangeEntry:
31 class PiecewiseBackend:
32 def __init__(
33 self,
29 - graph: fx.GraphModule,
34 + graph: fx.GraphModule | None,
35 vllm_config: VllmConfig,
36 piecewise_compile_index: int,
37 total_piecewise_compiles: int,
38 sym_shape_indices: list[int],
39 vllm_backend: VllmBackend,
40 + returns_tuple: bool,
41 + compiled_runnables: dict[str, Callable] | None = None,
42 ):
43 """
44 The backend for piecewise compilation.
@@ -41,13 +48,28 @@ class PiecewiseBackend:
48 We will compile `self.graph` once for the general shape,
49 and then compile for different shapes specified in
50 `compilation_config.compile_sizes`.
51 +
52 + This class supports two mutually exclusive modes:
53 + 1. Compilation (graph is set, compiled_runnables is None):
54 + Used during initial compilation when we have the FX graph
55 + and need to compile it for each shape range.
56 + 2. Precompilation (graph is None, compiled_runnables is set):
57 + Used when loading from cache/AOT artifacts where we already
58 + have pre-compiled callables and don't need the original graph.
59 +
60 + Exactly one of graph or compiled_runnables must be provided.
61 """
62 + assert bool(graph is not None) ^ bool(compiled_runnables is not None), (
63 + "exactly one of graph and compiled_runnables should be set."
64 + )
65 +
66 self.graph = graph
67 self.vllm_config = vllm_config
68 self.compilation_config = vllm_config.compilation_config
69 self.piecewise_compile_index = piecewise_compile_index
70 self.total_piecewise_compiles = total_piecewise_compiles
71 self.vllm_backend = vllm_backend
72 + self.compiled_runnables = compiled_runnables
73
74 self.is_first_graph = piecewise_compile_index == 0
75 self.is_last_graph = piecewise_compile_index == total_piecewise_compiles - 1
@@ -77,6 +99,7 @@ class PiecewiseBackend:
99 logger.debug_once(log_string)
100
101 self.sym_shape_indices = sym_shape_indices
102 + self.returns_tuple = returns_tuple
103
104 # the entries for ranges that we need to either
105 self.range_entries: dict[Range, RangeEntry] = {}
@@ -108,12 +131,71 @@ class PiecewiseBackend:
131 compile_range=range,
132 )
133
134 + # get the on_compilation_complete callback from context...
135 + # PiecewiseBackend is created during the first call,
136 + # which is when the context is set (see compilation/decorators.py)
137 + from vllm.compilation.backends import _on_compilation_complete_callback
138 +
139 + self.on_compilation_complete = _on_compilation_complete_callback.get()
140 +
141 + def get_compiled_graph_wrapper(self, compiled_graph):
142 + def compiled_graph_wrapper(*args):
143 + graph_output = compiled_graph(*args)
144 + # unpack the tuple if needed
145 + # TODO(rzou): the implication is that we're not
146 + # reading the python bytecode correctly in vLLM?
147 + if self.returns_tuple or not isinstance(graph_output, (tuple, list)):
148 + return graph_output
149 + else:
150 + return graph_output[0]
151 +
152 + return compiled_graph_wrapper
153 +
154 def check_for_ending_compilation(self) -> None:
155 if self.is_last_graph and not self.to_be_compiled_ranges:
156 # no specific sizes to compile
157 # save the hash of the inductor graph for the next run
158 self.vllm_backend.compiler_manager.save_to_file()
159 end_monitoring_torch_compile(self.vllm_config)
160 + # Call the completion callback (e.g., to save AOT compiled function)
161 + if self.on_compilation_complete is not None:
162 + self.on_compilation_complete()
163 +
164 + def to_bytes(self) -> dict[str, bytes]:
165 + class StandaloneCompiledArtifactsPickler(Pickler):
166 + def reducer_override(self, obj):
167 + if isinstance(obj, CachingAutotuner):
168 + obj.prepare_for_pickle()
169 + return pickle.loads, (
170 + pickle.dumps(
171 + obj,
172 + ),
173 + )
174 + return NotImplemented
175 +
176 + def serialize(fn) -> bytes:
177 + assert hasattr(fn, "serialize"), "fn must have serialize method"
178 + with torch._functorch.config.patch("bundled_autograd_cache", True):
179 + entry = fn.serialize()
180 +
181 + f = io.BytesIO()
182 + StandaloneCompiledArtifactsPickler(f).dump(entry)
183 + result = f.getvalue()
184 + return result
185 +
186 + out = {}
187 +
188 + for range_key, entry in self.range_entries.items():
189 + if not entry.compiled:
190 + logger.debug(
191 + "entry with range %s not compiled, so cannot get its bytes",
192 + range_key,
193 + )
194 + continue
195 + if hasattr(entry.runnable, "serialize"):
196 + out[str(range_key)] = serialize(entry.runnable)
197 +
198 + return out
199
200 def _fakify_args(self, args: tuple[Any, ...]) -> list[Any]:
201 # We need to pass fake example_inputs, otherwise torch.compile
@@ -127,6 +209,7 @@ class PiecewiseBackend:
209 # non fake tensors as example inputs!
210 # See issue https://github.com/vllm-project/vllm/issues/27899
211 fake_example_inputs = []
212 + assert self.graph is not None
213 for node in self.graph.graph.nodes:
214 # All place holders come first
215 if node.op == "placeholder":
@@ -140,28 +223,37 @@ class PiecewiseBackend:
223 self, range_entry: RangeEntry, args: tuple[Any, ...]
224 ) -> Any:
225 if not range_entry.compiled:
226 + if self.compiled_runnables is not None:
227 + range_entry.runnable = self.get_compiled_graph_wrapper(
228 + self.compiled_runnables[str(range_entry.compile_range)]
229 + )
230 + else:
231 + # args are real arguments
232 + # fakify for range, real args for concrete size.
233 + # For concrete size, we clear the shape env in
234 + # compiler_manager.compile() so no need to fakify.
235 + args_list = (
236 + self._fakify_args(args)
237 + if not range_entry.compile_range.is_single_size()
238 + else list(args)
239 + )
240 +
241 + with (
242 + torch._functorch.config.patch("bundled_autograd_cache", True),
243 + ):
244 + range_entry.runnable = self.vllm_backend.compiler_manager.compile(
245 + self.graph,
246 + args_list,
247 + self.vllm_backend.inductor_config,
248 + self.compilation_config,
249 + compile_range=range_entry.compile_range,
250 + graph_index=self.piecewise_compile_index,
251 + num_graphs=self.total_piecewise_compiles,
252 + )
253 +
254 range_entry.compiled = True
255 self.to_be_compiled_ranges.remove(range_entry.compile_range)
256
146 - # args are real arguments
147 - # fakify for range, real args for concrete size.
148 - # For concrete size, we clear the shape env in
149 - # compiler_manager.compile() so no need to fakify.
150 - args_list = (
151 - self._fakify_args(args)
152 - if not range_entry.compile_range.is_single_size()
153 - else list(args)
154 - )
155 - range_entry.runnable = self.vllm_backend.compiler_manager.compile(
156 - self.graph,
157 - args_list,
158 - self.vllm_backend.inductor_config,
159 - self.compilation_config,
160 - compile_range=range_entry.compile_range,
161 - graph_index=self.piecewise_compile_index,
162 - num_graphs=self.total_piecewise_compiles,
163 - )
164 -
257 self.check_for_ending_compilation()
258
259 def _find_range_for_shape(self, runtime_shape: int) -> RangeEntry | None:
vllm/envs.py
+8
@@ -108,6 +108,7 @@ if TYPE_CHECKING:
108 VLLM_USE_AOT_COMPILE: bool = False
109 VLLM_USE_BYTECODE_HOOK: bool = False
110 VLLM_FORCE_AOT_LOAD: bool = False
111 + VLLM_USE_MEGA_AOT_ARTIFACT: bool = False
112 VLLM_USE_TRITON_AWQ: bool = False
113 VLLM_ALLOW_RUNTIME_LORA_UPDATING: bool = False
114 VLLM_SKIP_P2P_CHECK: bool = False
@@ -630,6 +631,13 @@ environment_variables: dict[str, Callable[[], Any]] = {
631 # to load will result in a hard error when this is enabled.
632 # Will be ignored when VLLM_USE_AOT_COMPILE is disabled.
633 "VLLM_FORCE_AOT_LOAD": lambda: os.environ.get("VLLM_FORCE_AOT_LOAD", "0") == "1",
634 + # Enable loading compiled models directly from cached standalone compile artifacts
635 + # without re-splitting graph modules. This reduces overhead during model
636 + # loading by using reconstruct_serializable_fn_from_mega_artifact.
637 + "VLLM_USE_MEGA_AOT_ARTIFACT": lambda: os.environ.get(
638 + "VLLM_USE_MEGA_AOT_ARTIFACT", "0"
639 + )
640 + == "1",
641 # local rank of the process in the distributed setting, used to determine
642 # the GPU device id
643 "LOCAL_RANK": lambda: int(os.environ.get("LOCAL_RANK", "0")),