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 12, 14:45 UTC
Allow oot custom compiler extension via CompilerInterface (#28623)
Signed-off-by: wxsIcey <1790571317@qq.com> Signed-off-by: Mengqing Cao <cmq0113@163.com> Signed-off-by: Icey <1790571317@qq.com> Co-authored-by: Mengqing Cao <cmq0113@163.com>
Icey committed
Nov 25, 2025 at 15:25 UTC
888152bf87d62c9f5929d06f386068990b618db7
3 files changed
+42
-24
vllm/compilation/backends.py
+17
-17
@@ -63,13 +63,14 @@ def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
63
else:
64
logger.debug("Using InductorAdaptor")
65
return InductorAdaptor()
66
- else:
67
- assert compilation_config.backend == "eager", (
68
- "Custom backends not supported with CompilationMode.VLLM_COMPILE"
69
- )
70
-
66
+ elif compilation_config.backend == "eager":
67
logger.debug("Using EagerAdaptor")
68
return EagerAdaptor()
69
+ else:
70
+ logger.debug("Using custom backend: %s", compilation_config.backend)
71
+ compiler = resolve_obj_by_qualname(current_platform.get_compile_backend())()
72
+ assert isinstance(compiler, CompilerInterface)
73
+ return compiler
74
75
76
class CompilerManager:
@@ -545,7 +546,10 @@ class VllmBackend:
546
self.prefix = prefix or model_tag
547
548
# Passes to run on the graph post-grad.
548
- self.post_grad_pass_manager = PostGradPassManager()
549
+ self.pass_manager = resolve_obj_by_qualname(
550
+ current_platform.get_pass_manager_cls()
551
+ )()
552
+ self.pass_key = current_platform.pass_key
553
554
self.sym_tensor_indices = []
555
self.input_buffers = []
@@ -562,24 +566,20 @@ class VllmBackend:
566
567
def configure_post_pass(self):
568
config = self.compilation_config
565
- self.post_grad_pass_manager.configure(self.vllm_config)
569
+ self.pass_manager.configure(self.vllm_config)
570
571
# Post-grad custom passes are run using the post_grad_custom_post_pass
572
# hook. If a pass for that hook exists, add it to the pass manager.
573
inductor_config = config.inductor_compile_config
570
- PASS_KEY = "post_grad_custom_post_pass"
571
- if PASS_KEY in inductor_config:
572
- if isinstance(inductor_config[PASS_KEY], PostGradPassManager):
574
+ if self.pass_key in inductor_config:
575
+ if isinstance(inductor_config[self.pass_key], PostGradPassManager):
576
# PassManager already added to config, make sure it's correct
574
- assert (
575
- inductor_config[PASS_KEY].uuid()
576
- == self.post_grad_pass_manager.uuid()
577
- )
577
+ assert inductor_config[self.pass_key].uuid() == self.pass_manager.uuid()
578
else:
579
# Config should automatically wrap all inductor passes
580
- assert isinstance(inductor_config[PASS_KEY], InductorPass)
581
- self.post_grad_pass_manager.add(inductor_config[PASS_KEY])
582
- inductor_config[PASS_KEY] = self.post_grad_pass_manager
580
+ assert isinstance(inductor_config[self.pass_key], InductorPass)
581
+ self.pass_manager.add(inductor_config[self.pass_key])
582
+ inductor_config[self.pass_key] = self.pass_manager
583
584
def __call__(
585
self, graph: fx.GraphModule, example_inputs
vllm/config/compilation.py
+5
-7
@@ -331,9 +331,9 @@ class CompilationConfig:
331
We use string to avoid serialization issues when using compilation in a
332
distributed setting. When the compilation mode is 1 or 2, the backend is
333
used for the compilation directly (it sees the whole graph). When the
334
- compilation mode is 3, the backend is used for the piecewise compilation
335
- (it sees a part of the graph). The backend can not be custom for compilation
336
- mode 3, i.e. the backend must be either eager or inductor. Furthermore,
334
+ compilation mode is 3, the backend supports both whole graph and piecewise
335
+ compilation, available backends include eager, inductor, and custom backends,
336
+ the latter of which can be defined via `get_compile_backend`. Furthermore,
337
compilation is only piecewise if splitting ops is set accordingly and
338
use_inductor_graph_partition is off. Note that the default options for
339
splitting ops are sufficient for piecewise compilation.
@@ -768,7 +768,7 @@ class CompilationConfig:
768
self.backend = "inductor" if self.use_inductor else "eager"
769
770
if self.backend == "":
771
- self.backend = current_platform.simple_compile_backend
771
+ self.backend = current_platform.get_compile_backend()
772
773
def init_backend(self, vllm_config: "VllmConfig") -> str | Callable:
774
"""
@@ -800,9 +800,7 @@ class CompilationConfig:
800
801
assert self.mode == CompilationMode.VLLM_COMPILE
802
if self.backend not in ["eager", "inductor"]:
803
- raise ValueError(
804
- f"Invalid backend for piecewise compilation: {self.backend}"
805
- )
803
+ logger.info("Using OOT custom backend for compilation.")
804
805
from vllm.compilation.backends import VllmBackend
806
vllm/platforms/interface.py
+20
@@ -134,6 +134,11 @@ class Platform:
134
135
_global_graph_pool: Any | None = None
136
137
+ @property
138
+ def pass_key(self) -> str:
139
+ """Inductor config key for the PassManager custom pass"""
140
+ return "post_grad_custom_post_pass"
141
+
142
@property
143
def supported_dtypes(self) -> list[torch.dtype]:
144
"""Returns the supported dtypes for the current platform."""
@@ -177,6 +182,21 @@ class Platform:
182
# all ROCm platforms for now.
183
return self._enum in (PlatformEnum.CUDA, PlatformEnum.ROCM)
184
185
+ @classmethod
186
+ def get_pass_manager_cls(cls) -> str:
187
+ """
188
+ Get the pass manager class for this platform.
189
+ It will be registered as a custom pass under the current_platform.pass_key.
190
+ """
191
+ return "vllm.compilation.pass_manager.PostGradPassManager"
192
+
193
+ @classmethod
194
+ def get_compile_backend(cls) -> str:
195
+ """
196
+ Get the custom compile backend for current platform.
197
+ """
198
+ return cls.simple_compile_backend
199
+
200
@classmethod
201
def device_id_to_physical_device_id(cls, device_id: int):
202
# Treat empty device control env var as unset. This is a valid