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 14, 05:30 UTC
[torch.compile] Passing only necessary compilation config to inductor pass config (#27041)
Signed-off-by: Lu Fang <fanglu@fb.com> Co-authored-by: Lucia (Lu) Fang <fanglu@meta.com>
Lucia Fang committed
Oct 16, 2025 at 17:01 UTC
11ae016bd7836be0a17edd3ca3d9eae262a4f3a5
4 files changed
+47
-2
tests/compile/test_async_tp.py
+9
@@ -341,6 +341,15 @@ def async_tp_pass_on_test_model(
341
async_tp_pass = AsyncTPPass(vllm_config)
342
backend = TestBackend(async_tp_pass)
343
344
+ assert (
345
+ async_tp_pass.compilation_config.splitting_ops
346
+ == vllm_config.compilation_config.splitting_ops
347
+ )
348
+ assert (
349
+ async_tp_pass.compilation_config.use_inductor_graph_partition
350
+ == vllm_config.compilation_config.use_inductor_graph_partition
351
+ )
352
+
353
model = test_model_cls(hidden_size, dtype) # Pass dtype to model constructor
354
355
hidden_states = torch.randn(
tests/compile/test_config.py
+17
@@ -1,8 +1,11 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+import copy
4
+
5
import pytest
6
7
from vllm.compilation.counter import compilation_counter
8
+from vllm.compilation.fix_functionalization import FixFunctionalizationPass
9
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
10
from vllm.config.compilation import CompilationMode
11
from vllm.utils import _is_torch_equal_or_newer, is_torch_equal_or_newer
@@ -25,6 +28,20 @@ def test_use_cudagraphs_dynamic():
28
assert vllm_config.compilation_config.use_cudagraph
29
30
31
+def test_copy_pass():
32
+ vllm_config = VllmConfig()
33
+ inductor_pass = FixFunctionalizationPass(vllm_config)
34
+ copied_inductor_pass = copy.deepcopy(inductor_pass)
35
+ assert (
36
+ copied_inductor_pass.compilation_config.use_inductor_graph_partition
37
+ == vllm_config.compilation_config.use_inductor_graph_partition
38
+ )
39
+ assert (
40
+ copied_inductor_pass.compilation_config.splitting_ops
41
+ == vllm_config.compilation_config.splitting_ops
42
+ )
43
+
44
+
45
def test_custom_op():
46
# proper syntax
47
_ = CompilationConfig(custom_ops=["+quant_fp8", "-silu_and_mul"])
tests/compile/test_sequence_parallelism.py
+8
@@ -285,6 +285,14 @@ def sequence_parallelism_pass_on_test_model(
285
286
noop_pass = NoOpEliminationPass(vllm_config)
287
sequence_parallelism_pass = SequenceParallelismPass(vllm_config)
288
+ assert (
289
+ sequence_parallelism_pass.compilation_config.splitting_ops
290
+ == vllm_config.compilation_config.splitting_ops
291
+ )
292
+ assert (
293
+ sequence_parallelism_pass.compilation_config.use_inductor_graph_partition
294
+ == vllm_config.compilation_config.use_inductor_graph_partition
295
+ )
296
func_pass = FixFunctionalizationPass(vllm_config)
297
cleanup_pass = PostCleanupPass(vllm_config)
298
vllm/compilation/vllm_inductor_pass.py
+13
-2
@@ -3,7 +3,7 @@
3
import functools
4
import operator
5
import time
6
-import weakref
6
+from dataclasses import dataclass
7
from typing import ClassVar
8
9
import regex as re
@@ -19,6 +19,12 @@ from .inductor_pass import InductorPass
19
logger = init_logger(__name__)
20
21
22
+@dataclass
23
+class InductorCompilationConfig:
24
+ splitting_ops: list[str] | None = None
25
+ use_inductor_graph_partition: bool = False
26
+
27
+
28
class VllmInductorPass(InductorPass):
29
"""
30
An inductor pass with access to vLLM PassConfig.
@@ -29,7 +35,12 @@ class VllmInductorPass(InductorPass):
35
"""Keep track of pass index for debug dump ordering."""
36
37
def __init__(self, config: VllmConfig):
32
- self.compilation_config = weakref.proxy(config.compilation_config)
38
+ # Get only the necessary CompilationConfig for the inductor pass, since
39
+ # full `CompilationConfig` contains pointer to model which is unsafe.
40
+ self.compilation_config = InductorCompilationConfig(
41
+ splitting_ops=config.compilation_config.splitting_ops,
42
+ use_inductor_graph_partition=config.compilation_config.use_inductor_graph_partition,
43
+ )
44
self.pass_config = config.compilation_config.pass_config
45
self.model_dtype = config.model_config.dtype if config.model_config else None
46
self.device = config.device_config.device if config.device_config else None