backends.py 29.9 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
import ast
4
import dataclasses
5
6
import os
import pprint
7
import time
8
from contextlib import ExitStack
9
from typing import Any, Callable, Dict, List, Optional, Sequence, Set, Tuple
10
from unittest.mock import patch
11
12
13
14

import torch
import torch.fx as fx

15
import vllm.envs as envs
16
from vllm.config import CompilationConfig, VllmConfig
17
from vllm.logger import init_logger
18
from vllm.utils import weak_ref_tensors
19

20
21
from .compiler_interface import (CompilerInterface, EagerAdaptor,
                                 InductorAdaptor, InductorStandaloneAdaptor)
22
from .counter import compilation_counter
23
from .inductor_pass import InductorPass
24
from .monitor import end_monitoring_torch_compile
25
from .pass_manager import PostGradPassManager
26
27
28

logger = init_logger(__name__)

29

30
31
32
33
34
35
36
37
38
39
40
41
42
def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
    if compilation_config.use_inductor:
        if envs.VLLM_TEST_STANDALONE_COMPILE:
            logger.info("Using InductorStandaloneAdaptor")
            return InductorStandaloneAdaptor()
        else:
            logger.info("Using InductorAdaptor")
            return InductorAdaptor()
    else:
        logger.info("Using EagerAdaptor")
        return EagerAdaptor()


43
44
45
46
47
class CompilerManager:
    """
    A manager to manage the compilation process, including
    caching the compiled graph, loading the compiled graph,
    and compiling the graph.
48

49
50
51
    The cache is a dict mapping
    `(runtime_shape, graph_index, backend_name)`
    to `any_data` returned from the compiler.
52

53
54
55
    When serializing the cache, we save it to a Python file
    for readability. We don't use json here because json doesn't
    support int as key.
56
57
    """

58
    def __init__(self, compilation_config: CompilationConfig):
59
        self.cache: Dict[Tuple[Optional[int], int, str], Any] = dict()
60
        self.is_cache_updated = False
61
62
        self.compilation_config = compilation_config
        self.compiler = make_compiler(compilation_config)
63

64
65
    def compute_hash(self, vllm_config: VllmConfig) -> str:
        return self.compiler.compute_hash(vllm_config)
66

67
68
    def initialize_cache(self, cache_dir: str, disable_cache: bool = False):
        self.disable_cache = disable_cache
69
        self.cache_dir = cache_dir
70
71
72
73
        self.cache_file_path = os.path.join(cache_dir, "vllm_compile_cache.py")

        if not disable_cache and os.path.exists(self.cache_file_path):
            # load the cache from the file
74
            with open(self.cache_file_path) as f:
75
76
77
78
79
80
81
                # we use ast.literal_eval to parse the data
                # because it is a safe way to parse Python literals.
                # do not use eval(), it is unsafe.
                self.cache = ast.literal_eval(f.read())

        self.compiler.initialize_cache(cache_dir=cache_dir,
                                       disable_cache=disable_cache)
82
83

    def save_to_file(self):
84
        if self.disable_cache or not self.is_cache_updated:
85
            return
86
87
        printer = pprint.PrettyPrinter(indent=4)
        data = printer.pformat(self.cache)
88
        with open(self.cache_file_path, "w") as f:
89
90
91
92
93
94
95
96
97
98
99
100
            f.write(data)

    def load(self,
             graph: fx.GraphModule,
             example_inputs: List[Any],
             graph_index: int,
             runtime_shape: Optional[int] = None) -> Optional[Callable]:
        if (runtime_shape, graph_index, self.compiler.name) not in self.cache:
            return None
        handle = self.cache[(runtime_shape, graph_index, self.compiler.name)]
        compiled_graph = self.compiler.load(handle, graph, example_inputs,
                                            graph_index, runtime_shape)
101
        logger.debug(
102
103
104
105
106
107
108
109
110
111
112
113
114
            "Directly load the %s-th graph for shape %s from %s via "
            "handle %s", graph_index, str(runtime_shape), self.compiler.name,
            handle)
        return compiled_graph

    def compile(self,
                graph: fx.GraphModule,
                example_inputs,
                additional_inductor_config,
                compilation_config: CompilationConfig,
                graph_index: int = 0,
                num_graphs: int = 1,
                runtime_shape: Optional[int] = None) -> Any:
115
        if graph_index == 0:
116
117
118
119
120
121
122
123
124
125
126
127
            # before compiling the first graph, record the start time
            global compilation_start_time
            compilation_start_time = time.time()

        compilation_counter.num_backend_compilations += 1

        compiled_graph = None

        # try to load from the cache
        compiled_graph = self.load(graph, example_inputs, graph_index,
                                   runtime_shape)
        if compiled_graph is not None:
128
129
130
131
132
133
134
135
            if graph_index == num_graphs - 1:
                # after loading the last graph for this shape, record the time.
                # there can be multiple graphs due to piecewise compilation.
                now = time.time()
                elapsed = now - compilation_start_time
                logger.info(
                    "Directly load the compiled graph(s) for shape %s "
                    "from the cache, took %.3f s", str(runtime_shape), elapsed)
136
137
138
139
            return compiled_graph

        # no compiler cached the graph, or the cache is disabled,
        # we need to compile it
140
141
142
143
144
145
        if isinstance(self.compiler, InductorAdaptor):
            # Let compile_fx generate a key for us
            maybe_key = None
        else:
            maybe_key = \
                f"artifact_shape_{runtime_shape}_subgraph_{graph_index}"
146
        compiled_graph, handle = self.compiler.compile(
147
148
            graph, example_inputs, additional_inductor_config, runtime_shape,
            maybe_key)
149
150
151
152
153
154
155

        assert compiled_graph is not None, "Failed to compile the graph"

        # store the artifact in the cache
        if handle is not None:
            self.cache[(runtime_shape, graph_index,
                        self.compiler.name)] = handle
156
            self.is_cache_updated = True
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
            if graph_index == 0:
                # adds some info logging for the first graph
                logger.info("Cache the graph of shape %s for later use",
                            str(runtime_shape))
            logger.debug(
                "store the %s-th graph for shape %s from %s via handle %s",
                graph_index, str(runtime_shape), self.compiler.name, handle)

        # after compiling the last graph, record the end time
        if graph_index == num_graphs - 1:
            now = time.time()
            elapsed = now - compilation_start_time
            compilation_config.compilation_time += elapsed
            if runtime_shape is None:
                logger.info("Compiling a graph for general shape takes %.2f s",
                            elapsed)
            else:
                logger.info("Compiling a graph for shape %s takes %.2f s",
                            runtime_shape, elapsed)
176

177
        return compiled_graph
178
179


180
181
182
@dataclasses.dataclass
class SplitItem:
    submod_name: str
183
    graph_id: int
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
    is_splitting_graph: bool
    graph: fx.GraphModule


def split_graph(graph: fx.GraphModule,
                ops: List[str]) -> Tuple[fx.GraphModule, List[SplitItem]]:
    # split graph by ops
    subgraph_id = 0
    node_to_subgraph_id = {}
    split_op_graphs = []
    for node in graph.graph.nodes:
        if node.op in ("output", "placeholder"):
            continue
        if node.op == 'call_function' and str(node.target) in ops:
            subgraph_id += 1
            node_to_subgraph_id[node] = subgraph_id
            split_op_graphs.append(subgraph_id)
            subgraph_id += 1
        else:
            node_to_subgraph_id[node] = subgraph_id

    # `keep_original_order` is important!
    # otherwise pytorch might reorder the nodes and
    # the semantics of the graph will change when we
    # have mutations in the graph
    split_gm = torch.fx.passes.split_module.split_module(
210
        graph,
211
212
213
        None,
        lambda node: node_to_subgraph_id[node],
        keep_original_order=True)
214

215
    outputs = []
216

217
    names = [name for (name, module) in split_gm.named_modules()]
218

219
220
221
222
    for name in names:
        if "." in name or name == "":
            # recursive child module or the root module
            continue
223

224
        module = getattr(split_gm, name)
225

226
        graph_id = int(name.replace("submod_", ""))
227
228
229
230
231
        outputs.append(
            SplitItem(name, graph_id, (graph_id in split_op_graphs), module))

    # sort by intetger graph_id, rather than string name
    outputs.sort(key=lambda x: x.graph_id)
232

233
    return split_gm, outputs
234
235


236
237
238
# we share the global graph pool among all the backends
global_graph_pool = None

239
240
compilation_start_time = 0.0

241
242
243
244
245
246

class PiecewiseCompileInterpreter(torch.fx.Interpreter):
    """Code adapted from `torch.fx.passes.shape_prop.ShapeProp`.
    It runs the given graph with fake inputs, and compile some
    submodules specified by `compile_submod_names` with the given
    compilation configs.
247
248
249
250
251

    NOTE: the order in `compile_submod_names` matters, because
    it will be used to determine the order of the compiled piecewise
    graphs. The first graph will handle logging, and the last graph
    has some special cudagraph output handling.
252
253
254
    """

    def __init__(self, module: torch.fx.GraphModule,
255
                 compile_submod_names: List[str], vllm_config: VllmConfig,
256
                 graph_pool, vllm_backend: "VllmBackend"):
257
258
259
260
        super().__init__(module)
        from torch._guards import detect_fake_mode
        self.fake_mode = detect_fake_mode()
        self.compile_submod_names = compile_submod_names
261
        self.compilation_config = vllm_config.compilation_config
262
        self.graph_pool = graph_pool
263
        self.vllm_config = vllm_config
264
        self.vllm_backend = vllm_backend
265
266
        # When True, it annoyingly dumps the torch.fx.Graph on errors.
        self.extra_traceback = False
267
268
269
270
271
272

    def run(self, *args):
        fake_args = [
            self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
            for t in args
        ]
273
274
        with self.fake_mode:
            return super().run(*fake_args)
275
276
277
278
279
280
281
282

    def call_module(self, target: torch.fx.node.Target,
                    args: Tuple[torch.fx.node.Argument,
                                ...], kwargs: Dict[str, Any]) -> Any:
        assert isinstance(target, str)
        output = super().call_module(target, args, kwargs)

        if target in self.compile_submod_names:
283
            index = self.compile_submod_names.index(target)
284
285
286
287
            submod = self.fetch_attr(target)
            sym_shape_indices = [
                i for i, x in enumerate(args) if isinstance(x, torch.SymInt)
            ]
288
            global compilation_start_time
289
290
            compiled_graph_for_general_shape = self.vllm_backend.\
                compiler_manager.compile(
291
292
                submod,
                args,
293
294
                self.compilation_config.inductor_compile_config,
                self.compilation_config,
295
296
                graph_index=index,
                num_graphs=len(self.compile_submod_names),
297
                runtime_shape=None)
298
299

            self.module.__dict__[target] = PiecewiseBackend(
300
                submod, self.vllm_config, self.graph_pool, index,
301
                len(self.compile_submod_names), sym_shape_indices,
302
                compiled_graph_for_general_shape, self.vllm_backend)
303
304
305
306
307
308

            compilation_counter.num_piecewise_capturable_graphs_seen += 1

        return output


309
class VllmBackend:
310
    """The compilation backend for `torch.compile` with vLLM.
311
312
    It is used for compilation level of `CompilationLevel.PIECEWISE`,
    where we customize the compilation.
313

314
315
    The major work of this backend is to split the graph into
    piecewise graphs, and pass them to the piecewise backend.
316

317
318
    This backend also adds the PostGradPassManager to Inductor config,
    which handles the post-grad passes.
319
    """
320

321
322
    vllm_config: VllmConfig
    compilation_config: CompilationConfig
323
324
325
326
327
328
329
330
    graph_pool: Any
    _called: bool = False
    # the graph we compiled
    graph: fx.GraphModule
    # the stiching graph module for all the piecewise graphs
    split_gm: fx.GraphModule
    piecewise_graphs: List[SplitItem]
    returned_callable: Callable
331
332
    # Inductor passes to run on the graph pre-defunctionalization
    post_grad_passes: Sequence[Callable]
333
334
    sym_tensor_indices: List[int]
    input_buffers: List[torch.Tensor]
335
    compiler_manager: CompilerManager
336

337
338
    def __init__(
        self,
339
        vllm_config: VllmConfig,
340
    ):
341
342
343
344
345
346
347
348
        global global_graph_pool
        if global_graph_pool is None:
            global_graph_pool = torch.cuda.graph_pool_handle()

        # TODO: in the future, if we want to use multiple
        # streams, it might not be safe to share a global pool.
        # only investigate this when we use multiple streams
        self.graph_pool = global_graph_pool
349
350
351

        # Passes to run on the graph post-grad.
        self.post_grad_pass_manager = PostGradPassManager()
352

353
354
355
        self.sym_tensor_indices = []
        self.input_buffers = []

356
357
        self.vllm_config = vllm_config
        self.compilation_config = vllm_config.compilation_config
358

359
        self.compiler_manager: CompilerManager = CompilerManager(
360
            self.compilation_config)
361

362
363
364
        # `torch.compile` is JIT compiled, so we don't need to
        # do anything here

365
    def configure_post_pass(self):
366
        config = self.compilation_config
367
        self.post_grad_pass_manager.configure(self.vllm_config)
368

369
370
        # Post-grad custom passes are run using the post_grad_custom_post_pass
        # hook. If a pass for that hook exists, add it to the pass manager.
371
        inductor_config = config.inductor_compile_config
372
373
374
        PASS_KEY = "post_grad_custom_post_pass"
        if PASS_KEY in inductor_config:
            # Config should automatically wrap all inductor passes
375
376
377
378
379
380
            if isinstance(inductor_config[PASS_KEY], PostGradPassManager):
                assert (inductor_config[PASS_KEY].uuid() ==
                        self.post_grad_pass_manager.uuid())
            else:
                assert isinstance(inductor_config[PASS_KEY], InductorPass)
                self.post_grad_pass_manager.add(inductor_config[PASS_KEY])
381
        inductor_config[PASS_KEY] = self.post_grad_pass_manager
382

383
384
    def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:

385
        vllm_config = self.vllm_config
386
387
388
389
390
391
        if not self.compilation_config.cache_dir:
            # no provided cache dir, generate one based on the known factors
            # that affects the compilation. if none of the factors change,
            # the cache dir will be the same so that we can reuse the compiled
            # graph.

392
            factors = []
393
394
395
396
397
            # 0. factors come from the env, for example, The values of
            # VLLM_PP_LAYER_PARTITION will affects the computation graph.
            env_hash = envs.compute_hash()
            factors.append(env_hash)

398
399
400
            # 1. factors come from the vllm_config (it mainly summarizes how the
            #    model is created)
            config_hash = vllm_config.compute_hash()
401
            factors.append(config_hash)
402
403
404
405
406
407
408
409
410
411
412
413

            # 2. factors come from the code files that are traced by Dynamo (
            #    it mainly summarizes how the model is used in forward pass)
            forward_code_files = list(
                sorted(self.compilation_config.traced_files))
            self.compilation_config.traced_files.clear()
            logger.debug(
                "Traced files (to be considered for compilation cache):\n%s",
                "\n".join(forward_code_files))
            hash_content = []
            for filepath in forward_code_files:
                hash_content.append(filepath)
414
415
416
417
                if filepath == "<string>":
                    # This means the function was dynamically generated, with
                    # e.g. exec(). We can't actually check these.
                    continue
418
419
420
                with open(filepath) as f:
                    hash_content.append(f.read())
            import hashlib
421
422
            code_hash = hashlib.md5("\n".join(hash_content).encode(),
                                    usedforsecurity=False).hexdigest()
423
424
425
426
427
428
429
            factors.append(code_hash)

            # 3. compiler hash
            compiler_hash = self.compiler_manager.compute_hash(vllm_config)
            factors.append(compiler_hash)

            # combine all factors to generate the cache dir
430
431
            hash_key = hashlib.md5(str(factors).encode(),
                                   usedforsecurity=False).hexdigest()[:10]
432
433

            cache_dir = os.path.join(
434
435
436
437
438
439
                envs.VLLM_CACHE_ROOT,
                "torch_compile_cache",
                hash_key,
            )
            self.compilation_config.cache_dir = cache_dir

440
441
442
443
444
        if compilation_counter.num_graphs_seen > 0:
            cache_dir = self.compilation_config.cache_dir + \
                f'-{compilation_counter.num_graphs_seen}'
        else:
            cache_dir = self.compilation_config.cache_dir
445
        os.makedirs(cache_dir, exist_ok=True)
446
        self.compilation_config.cache_dir = cache_dir
447
448
449
        rank = vllm_config.parallel_config.rank
        dp_rank = vllm_config.parallel_config.data_parallel_rank
        local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}")
450
        os.makedirs(local_cache_dir, exist_ok=True)
451
        self.compilation_config.local_cache_dir = local_cache_dir
452

453
454
455
        disable_cache = envs.VLLM_DISABLE_COMPILE_CACHE

        if disable_cache:
456
457
458
            logger.info("vLLM's torch.compile cache is disabled.")
        else:
            logger.info("Using cache directory: %s for vLLM's torch.compile",
459
                        local_cache_dir)
460

461
462
        self.compiler_manager.initialize_cache(local_cache_dir, disable_cache)

463
464
        # when dynamo calls the backend, it means the bytecode
        # transform and analysis are done
465
        compilation_counter.num_graphs_seen += 1
466
467
468
        from .monitor import torch_compile_start_time
        dynamo_time = time.time() - torch_compile_start_time
        logger.info("Dynamo bytecode transform time: %.2f s", dynamo_time)
469
        self.compilation_config.compilation_time += dynamo_time
470
471
472
473
474
475

        # we control the compilation process, each instance can only be
        # called once
        assert not self._called, "VllmBackend can only be called once"

        self.graph = graph
476
        self.configure_post_pass()
477
478

        self.split_gm, self.piecewise_graphs = split_graph(
479
            graph, self.compilation_config.splitting_ops)
480

481
        from torch._dynamo.utils import lazy_format_graph_code
482
483
484
485
486

        # depyf will hook lazy_format_graph_code and dump the graph
        # for debugging, no need to print the graph here
        lazy_format_graph_code("before split", self.graph)
        lazy_format_graph_code("after split", self.split_gm)
487

488
489
490
491
492
493
494
495
496
497
        compilation_counter.num_piecewise_graphs_seen += len(
            self.piecewise_graphs)
        submod_names_to_compile = [
            item.submod_name for item in self.piecewise_graphs
            if not item.is_splitting_graph
        ]

        # propagate the split graph to the piecewise backend,
        # compile submodules with symbolic shapes
        PiecewiseCompileInterpreter(self.split_gm, submod_names_to_compile,
498
499
                                    self.vllm_config, self.graph_pool,
                                    self).run(*example_inputs)
500

501
502
503
504
505
506
507
508
509
510
511
512
        graph_path = os.path.join(local_cache_dir, "computation_graph.py")
        if not os.path.exists(graph_path):
            # code adapted from https://github.com/thuml/depyf/blob/dab831108a752d1facc00acdd6d4243891845c37/depyf/explain/patched_lazy_format_graph_code.py#L30 # noqa
            # use `print_readable` because it can include submodules
            src = "from __future__ import annotations\nimport torch\n" + \
                self.split_gm.print_readable(print_output=False)
            src = src.replace("<lambda>", "GraphModule")
            with open(graph_path, "w") as f:
                f.write(src)

            logger.debug("Computation graph saved to %s", graph_path)

513
514
        self._called = True

515
516
        if not self.compilation_config.use_cudagraph or \
            not self.compilation_config.cudagraph_copy_inputs:
517
518
519
520
521
522
523
524
525
526
527
            return self.split_gm

        # if we need to copy input buffers for cudagraph
        from torch._guards import detect_fake_mode
        fake_mode = detect_fake_mode()
        fake_args = [
            fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
            for t in example_inputs
        ]

        # index of tensors that have symbolic shapes (batch size)
528
529
530
        # for weights and static buffers, they will have concrete shapes.
        # symbolic shape only happens for input tensors.
        from torch.fx.experimental.symbolic_shapes import is_symbolic
531
532
        self.sym_tensor_indices = [
            i for i, x in enumerate(fake_args)
533
534
            if isinstance(x, torch._subclasses.fake_tensor.FakeTensor) and \
                any(is_symbolic(d) for d in x.size())
535
536
537
538
539
540
541
542
543
        ]

        # compiler managed cudagraph input buffers
        # we assume the first run with symbolic shapes
        # has the maximum size among all the tensors
        self.input_buffers = [
            example_inputs[x].clone() for x in self.sym_tensor_indices
        ]

youkaichao's avatar
youkaichao committed
544
545
        # this is the callable we return to Dynamo to run
        def copy_and_call(*args):
546
547
548
549
550
551
552
553
554
555
556
557
558
559
            list_args = list(args)
            for i, index in enumerate(self.sym_tensor_indices):
                runtime_tensor = list_args[index]
                runtime_shape = runtime_tensor.shape[0]
                static_tensor = self.input_buffers[i][:runtime_shape]

                # copy the tensor to the static buffer
                static_tensor.copy_(runtime_tensor)

                # replace the tensor in the list_args to the static buffer
                list_args[index] = static_tensor
            return self.split_gm(*list_args)

        return copy_and_call
560
561
562
563
564
565


@dataclasses.dataclass
class ConcreteSizeEntry:
    runtime_shape: int
    need_to_compile: bool  # the size is in compile_sizes
566
    use_cudagraph: bool  # the size is in cudagraph_capture_sizes
567
568
569
570
571
572
573

    compiled: bool = False
    runnable: Callable = None  # type: ignore
    num_finished_warmup: int = 0
    cudagraph: Optional[torch.cuda.CUDAGraph] = None
    output: Optional[Any] = None

574
575
576
577
    # for cudagraph debugging, track the input addresses
    # during capture, and check if they are the same during replay
    input_addresses: Optional[List[int]] = None

578
579
580

class PiecewiseBackend:

581
582
583
    def __init__(self, graph: fx.GraphModule, vllm_config: VllmConfig,
                 graph_pool: Any, piecewise_compile_index: int,
                 total_piecewise_compiles: int, sym_shape_indices: List[int],
584
585
                 compiled_graph_for_general_shape: Callable,
                 vllm_backend: VllmBackend):
586
587
588
589
590
591
        """
        The backend for piecewise compilation.
        It mainly handles the compilation and cudagraph capturing.

        We will compile `self.graph` once for the general shape,
        and then compile for different shapes specified in
592
        `compilation_config.compile_sizes`.
593
594
595
596
597
598
599

        Independently, we will capture cudagraph for different shapes.

        If a shape needs both compilation and cudagraph, we will
        compile it first, and then capture cudagraph.
        """
        self.graph = graph
600
601
        self.vllm_config = vllm_config
        self.compilation_config = vllm_config.compilation_config
602
        self.graph_pool = graph_pool
603
604
        self.piecewise_compile_index = piecewise_compile_index
        self.total_piecewise_compiles = total_piecewise_compiles
605
        self.vllm_backend = vllm_backend
606
607
608
609

        self.is_first_graph = piecewise_compile_index == 0
        self.is_last_graph = (
            piecewise_compile_index == total_piecewise_compiles - 1)
610
611

        self.compile_sizes: Set[int] = set(
612
            self.compilation_config.compile_sizes)
613
614
        self.cudagraph_capture_sizes: Set[int] = set(
            self.compilation_config.cudagraph_capture_sizes
615
        ) if self.compilation_config.use_cudagraph else set()
616
617
618

        self.first_run_finished = False

619
        self.compiled_graph_for_general_shape = compiled_graph_for_general_shape  # noqa
620

621
        self.sym_shape_indices = sym_shape_indices
622

623
624
        self.is_debugging_mode = envs.VLLM_LOGGING_LEVEL == "DEBUG"

625
626
627
        # the entries for different shapes that we need to either
        # compile or capture cudagraph
        self.concrete_size_entries: Dict[int, ConcreteSizeEntry] = {}
628
629
630
631

        # to_be_compiled_sizes tracks the remaining sizes to compile,
        # and updates during the compilation process, so we need to copy it
        self.to_be_compiled_sizes: Set[int] = self.compile_sizes.copy()
632
        for shape in self.compile_sizes.union(self.cudagraph_capture_sizes):
633
634
635
            self.concrete_size_entries[shape] = ConcreteSizeEntry(
                runtime_shape=shape,
                need_to_compile=shape in self.compile_sizes,
636
                use_cudagraph=shape in self.cudagraph_capture_sizes,
637
638
            )

639
640
641
642
    def check_for_ending_compilation(self):
        if self.is_last_graph and not self.to_be_compiled_sizes:
            # no specific sizes to compile
            # save the hash of the inductor graph for the next run
643
            self.vllm_backend.compiler_manager.save_to_file()
644
645
            end_monitoring_torch_compile(self.vllm_config)

646
647
648
    def __call__(self, *args) -> Any:
        if not self.first_run_finished:
            self.first_run_finished = True
649
            self.check_for_ending_compilation()
650
651
652
653
654
655
656
657
            return self.compiled_graph_for_general_shape(*args)

        runtime_shape = args[self.sym_shape_indices[0]]
        if runtime_shape not in self.concrete_size_entries:
            # we don't need to do anything for this shape
            return self.compiled_graph_for_general_shape(*args)

        entry = self.concrete_size_entries[runtime_shape]
658

659
660
        if entry.runnable is None:
            entry.runnable = self.compiled_graph_for_general_shape
661

662
663
        if entry.need_to_compile and not entry.compiled:
            entry.compiled = True
664
            self.to_be_compiled_sizes.remove(runtime_shape)
665
            # args are real arguments
666
            entry.runnable = self.vllm_backend.compiler_manager.compile(
667
668
                self.graph,
                args,
669
670
                self.compilation_config.inductor_compile_config,
                self.compilation_config,
671
672
                graph_index=self.piecewise_compile_index,
                num_graphs=self.total_piecewise_compiles,
673
                runtime_shape=runtime_shape)
674

675
676
            # finished compilations for all required shapes
            if self.is_last_graph and not self.to_be_compiled_sizes:
677
                self.check_for_ending_compilation()
678

679
680
681
682
        if not entry.use_cudagraph:
            return entry.runnable(*args)

        if entry.cudagraph is None:
683
            if entry.num_finished_warmup < self.compilation_config.cudagraph_num_of_warmups:  # noqa
684
685
686
687
688
                entry.num_finished_warmup += 1
                if self.is_first_graph:
                    logger.debug(
                        "Warming up %s/%s for shape %s",
                        entry.num_finished_warmup,
689
                        self.compilation_config.cudagraph_num_of_warmups,
690
691
692
693
                        runtime_shape)
                return entry.runnable(*args)

            if self.is_first_graph:
694
695
696
697
698
                # Since we capture cudagraph for many different shapes and
                # capturing is fast, we don't need to log it for every shape.
                # We only log it in the debug mode.
                logger.debug("Capturing a cudagraph for shape %s",
                             runtime_shape)
699

700
701
702
703
            input_addresses = [
                x.data_ptr() for x in args if isinstance(x, torch.Tensor)
            ]
            entry.input_addresses = input_addresses
704
            cudagraph = torch.cuda.CUDAGraph()
705

706
707
708
709
710
711
712
713
714
715
716
717
718
719
720
721
722
723
724
725
726
727
728
            with ExitStack() as stack:
                if not self.is_first_graph:
                    # during every model forward, we will capture
                    # many pieces of cudagraphs (roughly one per layer).
                    # running gc again and again across layers will
                    # make the cudagraph capture very slow.
                    # therefore, we only run gc for the first graph,
                    # and disable gc for the rest of the graphs.
                    stack.enter_context(patch("gc.collect", lambda: None))
                    stack.enter_context(
                        patch("torch.cuda.empty_cache", lambda: None))

                # mind-exploding: carefully manage the reference and memory.
                with torch.cuda.graph(cudagraph, pool=self.graph_pool):
                    # `output` is managed by pytorch's cudagraph pool
                    output = entry.runnable(*args)
                    if self.is_last_graph:
                        # by converting it to weak ref,
                        # the original `output` will immediately be released
                        # to save memory. It is only safe to do this for
                        # the last graph, because the output of the last graph
                        # will not be used by any other cuda graph.
                        output = weak_ref_tensors(output)
729
730
731
732
733

            # here we always use weak ref for the output
            # to save memory
            entry.output = weak_ref_tensors(output)
            entry.cudagraph = cudagraph
734
735
736

            compilation_counter.num_cudagraph_caputured += 1

737
738
739
740
741
742
743
744
745
746
747
748
749
750
            # important: we need to return the output, rather than
            # the weak ref of the output, so that pytorch can correctly
            # manage the memory during cuda graph capture
            return output

        if self.is_debugging_mode:
            # check if the input addresses are the same
            new_input_addresses = [
                x.data_ptr() for x in args if isinstance(x, torch.Tensor)
            ]
            assert new_input_addresses == entry.input_addresses, (
                "Input addresses for cudagraphs are different during replay."
                f" Expected {entry.input_addresses}, got {new_input_addresses}"
            )
751
752
753

        entry.cudagraph.replay()
        return entry.output