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 15, 09:45 UTC
[Compile][Platform] Make PiecewiseBackend pluggable and extendable (#18076)
Signed-off-by: Mengqing Cao <cmq0113@163.com> Co-authored-by: youkaichao <youkaichao@gmail.com>
Mengqing Cao committed
May 23, 2025 at 03:11 UTC
f8d2cc5f553e74a6eb8b15cd29da16bddb019af1
6 files changed
+305
-200
vllm/compilation/backends.py
+6
-200
@@ -6,9 +6,7 @@ import os
6
import pprint
7
import time
8
from collections.abc import Sequence
9
-from contextlib import ExitStack
9
from typing import Any, Callable, Optional
11
-from unittest.mock import patch
10
11
import torch
12
import torch.fx as fx
@@ -16,13 +14,13 @@ import torch.fx as fx
14
import vllm.envs as envs
15
from vllm.config import CompilationConfig, VllmConfig
16
from vllm.logger import init_logger
19
-from vllm.utils import weak_ref_tensors
17
+from vllm.platforms import current_platform
18
+from vllm.utils import resolve_obj_by_qualname
19
20
from .compiler_interface import (CompilerInterface, EagerAdaptor,
21
InductorAdaptor, InductorStandaloneAdaptor)
22
from .counter import compilation_counter
23
from .inductor_pass import InductorPass
25
-from .monitor import end_monitoring_torch_compile
24
from .pass_manager import PostGradPassManager
25
26
logger = init_logger(__name__)
@@ -297,7 +295,9 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
295
num_graphs=len(self.compile_submod_names),
296
runtime_shape=None)
297
300
- self.module.__dict__[target] = PiecewiseBackend(
298
+ piecewise_backend = resolve_obj_by_qualname(
299
+ current_platform.get_piecewise_backend_cls())
300
+ self.module.__dict__[target] = piecewise_backend(
301
submod, self.vllm_config, self.graph_pool, index,
302
len(self.compile_submod_names), sym_shape_indices,
303
compiled_graph_for_general_shape, self.vllm_backend)
@@ -341,7 +341,7 @@ class VllmBackend:
341
):
342
global global_graph_pool
343
if global_graph_pool is None:
344
- global_graph_pool = torch.cuda.graph_pool_handle()
344
+ global_graph_pool = current_platform.graph_pool_handle()
345
346
# TODO: in the future, if we want to use multiple
347
# streams, it might not be safe to share a global pool.
@@ -558,197 +558,3 @@ class VllmBackend:
558
return self.split_gm(*list_args)
559
560
return copy_and_call
561
-
562
-
563
-@dataclasses.dataclass
564
-class ConcreteSizeEntry:
565
- runtime_shape: int
566
- need_to_compile: bool # the size is in compile_sizes
567
- use_cudagraph: bool # the size is in cudagraph_capture_sizes
568
-
569
- compiled: bool = False
570
- runnable: Callable = None # type: ignore
571
- num_finished_warmup: int = 0
572
- cudagraph: Optional[torch.cuda.CUDAGraph] = None
573
- output: Optional[Any] = None
574
-
575
- # for cudagraph debugging, track the input addresses
576
- # during capture, and check if they are the same during replay
577
- input_addresses: Optional[list[int]] = None
578
-
579
-
580
-class PiecewiseBackend:
581
-
582
- def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,
583
- graph_pool: Any, piecewise_compile_index: int,
584
- total_piecewise_compiles: int, sym_shape_indices: list[int],
585
- compiled_graph_for_general_shape: Callable,
586
- vllm_backend: VllmBackend):
587
- """
588
- The backend for piecewise compilation.
589
- It mainly handles the compilation and cudagraph capturing.
590
-
591
- We will compile `self.graph` once for the general shape,
592
- and then compile for different shapes specified in
593
- `compilation_config.compile_sizes`.
594
-
595
- Independently, we will capture cudagraph for different shapes.
596
-
597
- If a shape needs both compilation and cudagraph, we will
598
- compile it first, and then capture cudagraph.
599
- """
600
- self.graph = graph
601
- self.vllm_config = vllm_config
602
- self.compilation_config = vllm_config.compilation_config
603
- self.graph_pool = graph_pool
604
- self.piecewise_compile_index = piecewise_compile_index
605
- self.total_piecewise_compiles = total_piecewise_compiles
606
- self.vllm_backend = vllm_backend
607
-
608
- self.is_first_graph = piecewise_compile_index == 0
609
- self.is_last_graph = (
610
- piecewise_compile_index == total_piecewise_compiles - 1)
611
-
612
- self.compile_sizes: set[int] = set(
613
- self.compilation_config.compile_sizes)
614
- self.cudagraph_capture_sizes: set[int] = set(
615
- self.compilation_config.cudagraph_capture_sizes
616
- ) if self.compilation_config.use_cudagraph else set()
617
-
618
- self.first_run_finished = False
619
-
620
- self.compiled_graph_for_general_shape = compiled_graph_for_general_shape # noqa
621
-
622
- self.sym_shape_indices = sym_shape_indices
623
-
624
- self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
625
-
626
- # the entries for different shapes that we need to either
627
- # compile or capture cudagraph
628
- self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {}
629
-
630
- # to_be_compiled_sizes tracks the remaining sizes to compile,
631
- # and updates during the compilation process, so we need to copy it
632
- self.to_be_compiled_sizes: set[int] = self.compile_sizes.copy()
633
- for shape in self.compile_sizes.union(self.cudagraph_capture_sizes):
634
- self.concrete_size_entries[shape] = ConcreteSizeEntry(
635
- runtime_shape=shape,
636
- need_to_compile=shape in self.compile_sizes,
637
- use_cudagraph=shape in self.cudagraph_capture_sizes,
638
- )
639
-
640
- def check_for_ending_compilation(self):
641
- if self.is_last_graph and not self.to_be_compiled_sizes:
642
- # no specific sizes to compile
643
- # save the hash of the inductor graph for the next run
644
- self.vllm_backend.compiler_manager.save_to_file()
645
- end_monitoring_torch_compile(self.vllm_config)
646
-
647
- def __call__(self, *args) -> Any:
648
- if not self.first_run_finished:
649
- self.first_run_finished = True
650
- self.check_for_ending_compilation()
651
- return self.compiled_graph_for_general_shape(*args)
652
-
653
- runtime_shape = args[self.sym_shape_indices[0]]
654
- if runtime_shape not in self.concrete_size_entries:
655
- # we don't need to do anything for this shape
656
- return self.compiled_graph_for_general_shape(*args)
657
-
658
- entry = self.concrete_size_entries[runtime_shape]
659
-
660
- if entry.runnable is None:
661
- entry.runnable = self.compiled_graph_for_general_shape
662
-
663
- if entry.need_to_compile and not entry.compiled:
664
- entry.compiled = True
665
- self.to_be_compiled_sizes.remove(runtime_shape)
666
- # args are real arguments
667
- entry.runnable = self.vllm_backend.compiler_manager.compile(
668
- self.graph,
669
- args,
670
- self.compilation_config.inductor_compile_config,
671
- self.compilation_config,
672
- graph_index=self.piecewise_compile_index,
673
- num_graphs=self.total_piecewise_compiles,
674
- runtime_shape=runtime_shape)
675
-
676
- # finished compilations for all required shapes
677
- if self.is_last_graph and not self.to_be_compiled_sizes:
678
- self.check_for_ending_compilation()
679
-
680
- if not entry.use_cudagraph:
681
- return entry.runnable(*args)
682
-
683
- if entry.cudagraph is None:
684
- if entry.num_finished_warmup < self.compilation_config.cudagraph_num_of_warmups: # noqa
685
- entry.num_finished_warmup += 1
686
- if self.is_first_graph:
687
- logger.debug(
688
- "Warming up %s/%s for shape %s",
689
- entry.num_finished_warmup,
690
- self.compilation_config.cudagraph_num_of_warmups,
691
- runtime_shape)
692
- return entry.runnable(*args)
693
-
694
- if self.is_first_graph:
695
- # Since we capture cudagraph for many different shapes and
696
- # capturing is fast, we don't need to log it for every shape.
697
- # We only log it in the debug mode.
698
- logger.debug("Capturing a cudagraph for shape %s",
699
- runtime_shape)
700
-
701
- input_addresses = [
702
- x.data_ptr() for x in args if isinstance(x, torch.Tensor)
703
- ]
704
- entry.input_addresses = input_addresses
705
- cudagraph = torch.cuda.CUDAGraph()
706
-
707
- with ExitStack() as stack:
708
- if not self.is_first_graph:
709
- # during every model forward, we will capture
710
- # many pieces of cudagraphs (roughly one per layer).
711
- # running gc again and again across layers will
712
- # make the cudagraph capture very slow.
713
- # therefore, we only run gc for the first graph,
714
- # and disable gc for the rest of the graphs.
715
- stack.enter_context(patch("gc.collect", lambda: None))
716
- stack.enter_context(
717
- patch("torch.cuda.empty_cache", lambda: None))
718
-
719
- # mind-exploding: carefully manage the reference and memory.
720
- with torch.cuda.graph(cudagraph, pool=self.graph_pool):
721
- # `output` is managed by pytorch's cudagraph pool
722
- output = entry.runnable(*args)
723
- if self.is_last_graph:
724
- # by converting it to weak ref,
725
- # the original `output` will immediately be released
726
- # to save memory. It is only safe to do this for
727
- # the last graph, because the output of the last graph
728
- # will not be used by any other cuda graph.
729
- output = weak_ref_tensors(output)
730
-
731
- # here we always use weak ref for the output
732
- # to save memory
733
- entry.output = weak_ref_tensors(output)
734
- entry.cudagraph = cudagraph
735
-
736
- compilation_counter.num_cudagraph_caputured += 1
737
-
738
- # important: we need to return the output, rather than
739
- # the weak ref of the output, so that pytorch can correctly
740
- # manage the memory during cuda graph capture
741
- return output
742
-
743
- if self.is_debugging_mode:
744
- # check if the input addresses are the same
745
- new_input_addresses = [
746
- x.data_ptr() for x in args if isinstance(x, torch.Tensor)
747
- ]
748
- assert new_input_addresses == entry.input_addresses, (
749
- "Input addresses for cudagraphs are different during replay."
750
- f" Expected {entry.input_addresses}, got {new_input_addresses}"
751
- )
752
-
753
- entry.cudagraph.replay()
754
- return entry.output
vllm/compilation/base_piecewise_backend.py
new
+71
@@ -0,0 +1,71 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+from typing import Any, Callable, Protocol
4
+
5
+import torch.fx as fx
6
+
7
+from vllm.compilation.backends import VllmBackend
8
+from vllm.config import VllmConfig
9
+
10
+
11
+class AbstractPiecewiseBackend(Protocol):
12
+ """
13
+ PiecewiseBackend interface that allows platforms to extend
14
+ piecewise static graph.
15
+ """
16
+
17
+ def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,
18
+ graph_pool: Any, piecewise_compile_index: int,
19
+ total_piecewise_compiles: int, sym_shape_indices: list[int],
20
+ compiled_graph_for_general_shape: Callable,
21
+ vllm_backend: VllmBackend, **kwargs):
22
+ """
23
+ Initializes the PiecewiseBackend class with compilation and
24
+ execution-related configurations.
25
+
26
+ This class handles piecewise compilation, graph capturing,
27
+ and dispatching for specific input shapes.
28
+
29
+ Args:
30
+ graph (fx.GraphModule): The graph represented in fx.
31
+ vllm_config (VllmConfig): Global configuration for vLLM.
32
+ graph_pool (Any):
33
+ Graph memory pool handle, e.g.,
34
+ `torch.cuda.graph_pool_handle()`.
35
+ piecewise_compile_index (int):
36
+ Index of the current piecewise subgraph.
37
+ total_piecewise_compiles (int):
38
+ Total number of piecewise-compiled graphs.
39
+ sym_shape_indices (list[int]):
40
+ Indices of symbolic shape.
41
+ compiled_graph_for_general_shape (Callable):
42
+ Callable that executes the graph compiled for general shapes.
43
+ vllm_backend (VllmBackend):
44
+ Backend compiler that manages compilation and graph runtime
45
+ for vLLM.
46
+
47
+ Keyword Args:
48
+ kwargs: Additional keyword arguments reserved for future
49
+ extensions or custom platforms.
50
+ """
51
+ raise NotImplementedError
52
+
53
+ def __call__(self, *args) -> Any:
54
+ """Executes the compiled graph for given input args.
55
+
56
+ If this is the first invocation, executes the general compiled graph
57
+ and initiates the compilation process tracking. For subsequent calls,
58
+ dynamically dispatches execution to either a compiled graph or a static
59
+ graph based on the input shape.
60
+
61
+ Args:
62
+ *args: Variable length input arguments to be passed into the
63
+ graph. The symbolic shape is expected to be in position
64
+ `sym_shape_indices[0]`.
65
+
66
+ Returns:
67
+ Any: Output of the executed graph. This can be from the general
68
+ compiled graph, a specialized compiled version for the given shape,
69
+ or a replayed static graph.
70
+ """
71
+ raise NotImplementedError
vllm/compilation/cuda_piecewise_backend.py
new
+213
@@ -0,0 +1,213 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+
3
+import dataclasses
4
+from contextlib import ExitStack
5
+from typing import Any, Callable, Optional
6
+from unittest.mock import patch
7
+
8
+import torch
9
+import torch.fx as fx
10
+
11
+import vllm.envs as envs
12
+from vllm.compilation.backends import VllmBackend
13
+from vllm.compilation.counter import compilation_counter
14
+from vllm.compilation.monitor import end_monitoring_torch_compile
15
+from vllm.config import VllmConfig
16
+from vllm.logger import init_logger
17
+from vllm.utils import weak_ref_tensors
18
+
19
+logger = init_logger(__name__)
20
+
21
+
22
+@dataclasses.dataclass
23
+class ConcreteSizeEntry:
24
+ runtime_shape: int
25
+ need_to_compile: bool # the size is in compile_sizes
26
+ use_cudagraph: bool # the size is in cudagraph_capture_sizes
27
+
28
+ compiled: bool = False
29
+ runnable: Callable = None # type: ignore
30
+ num_finished_warmup: int = 0
31
+ cudagraph: Optional[torch.cuda.CUDAGraph] = None
32
+ output: Optional[Any] = None
33
+
34
+ # for cudagraph debugging, track the input addresses
35
+ # during capture, and check if they are the same during replay
36
+ input_addresses: Optional[list[int]] = None
37
+
38
+
39
+class CUDAPiecewiseBackend:
40
+
41
+ def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,
42
+ graph_pool: Any, piecewise_compile_index: int,
43
+ total_piecewise_compiles: int, sym_shape_indices: list[int],
44
+ compiled_graph_for_general_shape: Callable,
45
+ vllm_backend: VllmBackend):
46
+ """
47
+ The backend for piecewise compilation.
48
+ It mainly handles the compilation and cudagraph capturing.
49
+
50
+ We will compile `self.graph` once for the general shape,
51
+ and then compile for different shapes specified in
52
+ `compilation_config.compile_sizes`.
53
+
54
+ Independently, we will capture cudagraph for different shapes.
55
+
56
+ If a shape needs both compilation and cudagraph, we will
57
+ compile it first, and then capture cudagraph.
58
+ """
59
+ self.graph = graph
60
+ self.vllm_config = vllm_config
61
+ self.compilation_config = vllm_config.compilation_config
62
+ self.graph_pool = graph_pool
63
+ self.piecewise_compile_index = piecewise_compile_index
64
+ self.total_piecewise_compiles = total_piecewise_compiles
65
+ self.vllm_backend = vllm_backend
66
+
67
+ self.is_first_graph = piecewise_compile_index == 0
68
+ self.is_last_graph = (
69
+ piecewise_compile_index == total_piecewise_compiles - 1)
70
+
71
+ self.compile_sizes: set[int] = set(
72
+ self.compilation_config.compile_sizes)
73
+ self.cudagraph_capture_sizes: set[int] = set(
74
+ self.compilation_config.cudagraph_capture_sizes
75
+ ) if self.compilation_config.use_cudagraph else set()
76
+
77
+ self.first_run_finished = False
78
+
79
+ self.compiled_graph_for_general_shape = compiled_graph_for_general_shape # noqa
80
+
81
+ self.sym_shape_indices = sym_shape_indices
82
+
83
+ self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"
84
+
85
+ # the entries for different shapes that we need to either
86
+ # compile or capture cudagraph
87
+ self.concrete_size_entries: dict[int, ConcreteSizeEntry] = {}
88
+
89
+ # to_be_compiled_sizes tracks the remaining sizes to compile,
90
+ # and updates during the compilation process, so we need to copy it
91
+ self.to_be_compiled_sizes: set[int] = self.compile_sizes.copy()
92
+ for shape in self.compile_sizes.union(self.cudagraph_capture_sizes):
93
+ self.concrete_size_entries[shape] = ConcreteSizeEntry(
94
+ runtime_shape=shape,
95
+ need_to_compile=shape in self.compile_sizes,
96
+ use_cudagraph=shape in self.cudagraph_capture_sizes,
97
+ )
98
+
99
+ def check_for_ending_compilation(self):
100
+ if self.is_last_graph and not self.to_be_compiled_sizes:
101
+ # no specific sizes to compile
102
+ # save the hash of the inductor graph for the next run
103
+ self.vllm_backend.compiler_manager.save_to_file()
104
+ end_monitoring_torch_compile(self.vllm_config)
105
+
106
+ def __call__(self, *args) -> Any:
107
+ if not self.first_run_finished:
108
+ self.first_run_finished = True
109
+ self.check_for_ending_compilation()
110
+ return self.compiled_graph_for_general_shape(*args)
111
+
112
+ runtime_shape = args[self.sym_shape_indices[0]]
113
+ if runtime_shape not in self.concrete_size_entries:
114
+ # we don't need to do anything for this shape
115
+ return self.compiled_graph_for_general_shape(*args)
116
+
117
+ entry = self.concrete_size_entries[runtime_shape]
118
+
119
+ if entry.runnable is None:
120
+ entry.runnable = self.compiled_graph_for_general_shape
121
+
122
+ if entry.need_to_compile and not entry.compiled:
123
+ entry.compiled = True
124
+ self.to_be_compiled_sizes.remove(runtime_shape)
125
+ # args are real arguments
126
+ entry.runnable = self.vllm_backend.compiler_manager.compile(
127
+ self.graph,
128
+ args,
129
+ self.compilation_config.inductor_compile_config,
130
+ self.compilation_config,
131
+ graph_index=self.piecewise_compile_index,
132
+ num_graphs=self.total_piecewise_compiles,
133
+ runtime_shape=runtime_shape)
134
+
135
+ # finished compilations for all required shapes
136
+ if self.is_last_graph and not self.to_be_compiled_sizes:
137
+ self.check_for_ending_compilation()
138
+
139
+ if not entry.use_cudagraph:
140
+ return entry.runnable(*args)
141
+
142
+ if entry.cudagraph is None:
143
+ if entry.num_finished_warmup < self.compilation_config.cudagraph_num_of_warmups: # noqa
144
+ entry.num_finished_warmup += 1
145
+ if self.is_first_graph:
146
+ logger.debug(
147
+ "Warming up %s/%s for shape %s",
148
+ entry.num_finished_warmup,
149
+ self.compilation_config.cudagraph_num_of_warmups,
150
+ runtime_shape)
151
+ return entry.runnable(*args)
152
+
153
+ if self.is_first_graph:
154
+ # Since we capture cudagraph for many different shapes and
155
+ # capturing is fast, we don't need to log it for every shape.
156
+ # We only log it in the debug mode.
157
+ logger.debug("Capturing a cudagraph for shape %s",
158
+ runtime_shape)
159
+
160
+ input_addresses = [
161
+ x.data_ptr() for x in args if isinstance(x, torch.Tensor)
162
+ ]
163
+ entry.input_addresses = input_addresses
164
+ cudagraph = torch.cuda.CUDAGraph()
165
+
166
+ with ExitStack() as stack:
167
+ if not self.is_first_graph:
168
+ # during every model forward, we will capture
169
+ # many pieces of cudagraphs (roughly one per layer).
170
+ # running gc again and again across layers will
171
+ # make the cudagraph capture very slow.
172
+ # therefore, we only run gc for the first graph,
173
+ # and disable gc for the rest of the graphs.
174
+ stack.enter_context(patch("gc.collect", lambda: None))
175
+ stack.enter_context(
176
+ patch("torch.cuda.empty_cache", lambda: None))
177
+
178
+ # mind-exploding: carefully manage the reference and memory.
179
+ with torch.cuda.graph(cudagraph, pool=self.graph_pool):
180
+ # `output` is managed by pytorch's cudagraph pool
181
+ output = entry.runnable(*args)
182
+ if self.is_last_graph:
183
+ # by converting it to weak ref,
184
+ # the original `output` will immediately be released
185
+ # to save memory. It is only safe to do this for
186
+ # the last graph, because the output of the last graph
187
+ # will not be used by any other cuda graph.
188
+ output = weak_ref_tensors(output)
189
+
190
+ # here we always use weak ref for the output
191
+ # to save memory
192
+ entry.output = weak_ref_tensors(output)
193
+ entry.cudagraph = cudagraph
194
+
195
+ compilation_counter.num_cudagraph_caputured += 1
196
+
197
+ # important: we need to return the output, rather than
198
+ # the weak ref of the output, so that pytorch can correctly
199
+ # manage the memory during cuda graph capture
200
+ return output
201
+
202
+ if self.is_debugging_mode:
203
+ # check if the input addresses are the same
204
+ new_input_addresses = [
205
+ x.data_ptr() for x in args if isinstance(x, torch.Tensor)
206
+ ]
207
+ assert new_input_addresses == entry.input_addresses, (
208
+ "Input addresses for cudagraphs are different during replay."
209
+ f" Expected {entry.input_addresses}, got {new_input_addresses}"
210
+ )
211
+
212
+ entry.cudagraph.replay()
213
+ return entry.output
vllm/platforms/cuda.py
+4
@@ -311,6 +311,10 @@ class CudaPlatformBase(Platform):
311
def use_custom_allreduce(cls) -> bool:
312
return True
313
314
+ @classmethod
315
+ def get_piecewise_backend_cls(cls) -> str:
316
+ return "vllm.compilation.cuda_piecewise_backend.CUDAPiecewiseBackend" # noqa
317
+
318
319
# NVML utils
320
# Note that NVML is not affected by `CUDA_VISIBLE_DEVICES`,
vllm/platforms/interface.py
+7
@@ -478,6 +478,13 @@ class Platform:
478
"""
479
raise NotImplementedError
480
481
+ @classmethod
482
+ def get_piecewise_backend_cls(cls) -> str:
483
+ """
484
+ Get piecewise backend class for piecewise graph.
485
+ """
486
+ return "vllm.compilation.base_piecewise_backend.AbstractPiecewiseBackend" # noqa
487
+
488
489
class UnspecifiedPlatform(Platform):
490
_enum = PlatformEnum.UNSPECIFIED
vllm/platforms/rocm.py
+4
@@ -382,3 +382,7 @@ class RocmPlatform(Platform):
382
@classmethod
383
def is_navi(cls) -> bool:
384
return 'gfx1' in torch.cuda.get_device_properties(0).gcnArchName
385
+
386
+ @classmethod
387
+ def get_piecewise_backend_cls(cls) -> str:
388
+ return "vllm.compilation.cuda_piecewise_backend.CUDAPiecewiseBackend" # noqa