backends.py 26.6 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3

4
import ast
5
import dataclasses
6
import hashlib
7
8
import os
import pprint
9
import time
10
from collections.abc import Sequence
11
from contextlib import contextmanager
12
from typing import Any, Callable, Optional
13
14
15

import torch
import torch.fx as fx
16
from torch._dispatch.python import enable_python_dispatcher
17

18
import vllm.envs as envs
19
20
21
22
23
from vllm.compilation.inductor_pass import pass_context
from vllm.compilation.partition_rules import (
    inductor_partition_rule_context,
    resolve_defined_ops,
)
24
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
25
from vllm.logger import init_logger
26
from vllm.platforms import current_platform
27
from vllm.utils import is_torch_equal_or_newer, resolve_obj_by_qualname
28

29
from .caching import VllmSerializableFunction
30
31
32
33
34
35
from .compiler_interface import (
    CompilerInterface,
    EagerAdaptor,
    InductorAdaptor,
    InductorStandaloneAdaptor,
)
36
from .counter import compilation_counter
37
38
from .inductor_pass import InductorPass
from .pass_manager import PostGradPassManager
39
40
41

logger = init_logger(__name__)

42

43
def make_compiler(compilation_config: CompilationConfig) -> CompilerInterface:
44
    if compilation_config.use_inductor:
45
46
        # Use standalone compile only if requested, version is new enough,
        # and the symbol actually exists in this PyTorch build.
47
48
49
50
51
        if (
            envs.VLLM_USE_STANDALONE_COMPILE
            and is_torch_equal_or_newer("2.8.0.dev")
            and hasattr(torch._inductor, "standalone_compile")
        ):
52
            logger.debug("Using InductorStandaloneAdaptor")
53
54
            return InductorStandaloneAdaptor()
        else:
55
            logger.debug("Using InductorAdaptor")
56
57
            return InductorAdaptor()
    else:
58
        logger.debug("Using EagerAdaptor")
59
60
61
        return EagerAdaptor()


62
63
64
65
66
class CompilerManager:
    """
    A manager to manage the compilation process, including
    caching the compiled graph, loading the compiled graph,
    and compiling the graph.
67

68
69
70
    The cache is a dict mapping
    `(runtime_shape, graph_index, backend_name)`
    to `any_data` returned from the compiler.
71

72
73
74
    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.
75
76
    """

77
    def __init__(self, compilation_config: CompilationConfig):
78
        self.cache: dict[tuple[Optional[int], int, str], Any] = dict()
79
        self.is_cache_updated = False
80
81
        self.compilation_config = compilation_config
        self.compiler = make_compiler(compilation_config)
82

83
84
    def compute_hash(self, vllm_config: VllmConfig) -> str:
        return self.compiler.compute_hash(vllm_config)
85

86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
    @contextmanager
    def compile_context(self, runtime_shape: Optional[int] = None):
        """Provide compilation context for the duration of compilation to set
        any torch global properties we want to scope to a single Inductor
        compilation (e.g. partition rules, pass context)."""
        with pass_context(runtime_shape):
            if self.compilation_config.use_inductor_graph_partition:
                inductor_partition_ops = resolve_defined_ops(
                    self.compilation_config.splitting_ops
                )
                with inductor_partition_rule_context(inductor_partition_ops):
                    yield
            else:
                yield

101
102
103
    def initialize_cache(
        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
    ):
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
        """
        Initialize the cache directory for the compiler.

        The organization of the cache directory is as follows:
        cache_dir=/path/to/hash_str/rank_i_j/prefix/
        inside cache_dir, there will be:
        - vllm_compile_cache.py
        - computation_graph.py
        - transformed_code.py

        for multiple prefixes, they can share the same
        base cache dir of /path/to/hash_str/rank_i_j/ ,
        to store some common compilation artifacts.
        """

119
        self.disable_cache = disable_cache
120
        self.cache_dir = cache_dir
121
122
123
124
        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
125
            with open(self.cache_file_path) as f:
126
127
128
129
130
                # 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())

131
132
133
        self.compiler.initialize_cache(
            cache_dir=cache_dir, disable_cache=disable_cache, prefix=prefix
        )
134
135

    def save_to_file(self):
136
        if self.disable_cache or not self.is_cache_updated:
137
            return
138
139
        printer = pprint.PrettyPrinter(indent=4)
        data = printer.pformat(self.cache)
140
        with open(self.cache_file_path, "w") as f:
141
142
            f.write(data)

143
144
145
146
147
148
149
    def load(
        self,
        graph: fx.GraphModule,
        example_inputs: list[Any],
        graph_index: int,
        runtime_shape: Optional[int] = None,
    ) -> Optional[Callable]:
150
151
152
        if (runtime_shape, graph_index, self.compiler.name) not in self.cache:
            return None
        handle = self.cache[(runtime_shape, graph_index, self.compiler.name)]
153
154
155
        compiled_graph = self.compiler.load(
            handle, graph, example_inputs, graph_index, runtime_shape
        )
156
157
        if runtime_shape is None:
            logger.debug(
158
159
160
161
162
                "Directly load the %s-th graph for dynamic shape from %s via handle %s",
                graph_index,
                self.compiler.name,
                handle,
            )
163
164
        else:
            logger.debug(
165
166
167
168
169
170
                "Directly load the %s-th graph for shape %s from %s via handle %s",
                graph_index,
                str(runtime_shape),
                self.compiler.name,
                handle,
            )
171
172
        return compiled_graph

173
174
175
176
177
178
179
180
181
182
    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:
183
        if graph_index == 0:
184
185
186
187
188
189
190
191
192
            # 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
193
        compiled_graph = self.load(graph, example_inputs, graph_index, runtime_shape)
194
        if compiled_graph is not None:
195
196
197
198
199
            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
200
                compilation_config.compilation_time += elapsed
201
202
203
                if runtime_shape is None:
                    logger.info(
                        "Directly load the compiled graph(s) for dynamic shape "
204
205
206
                        "from the cache, took %.3f s",
                        elapsed,
                    )
207
208
209
                else:
                    logger.info(
                        "Directly load the compiled graph(s) for shape %s "
210
211
212
213
                        "from the cache, took %.3f s",
                        str(runtime_shape),
                        elapsed,
                    )
214
215
216
217
            return compiled_graph

        # no compiler cached the graph, or the cache is disabled,
        # we need to compile it
218
219
220
221
        if isinstance(self.compiler, InductorAdaptor):
            # Let compile_fx generate a key for us
            maybe_key = None
        else:
222
            maybe_key = f"artifact_shape_{runtime_shape}_subgraph_{graph_index}"
223
224
225
226
227
228
229
230
231

        with self.compile_context(runtime_shape):
            compiled_graph, handle = self.compiler.compile(
                graph,
                example_inputs,
                additional_inductor_config,
                runtime_shape,
                maybe_key,
            )
232
233
234
235

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

        # store the artifact in the cache
236
        if not envs.VLLM_DISABLE_COMPILE_CACHE and handle is not None:
237
            self.cache[(runtime_shape, graph_index, self.compiler.name)] = handle
238
            compilation_counter.num_cache_entries_updated += 1
239
            self.is_cache_updated = True
240
241
            if graph_index == 0:
                # adds some info logging for the first graph
242
                if runtime_shape is None:
243
                    logger.info("Cache the graph for dynamic shape for later use")
244
                else:
245
246
247
                    logger.info(
                        "Cache the graph of shape %s for later use", str(runtime_shape)
                    )
248
249
            if runtime_shape is None:
                logger.debug(
250
251
252
253
254
                    "Store the %s-th graph for dynamic shape from %s via handle %s",
                    graph_index,
                    self.compiler.name,
                    handle,
                )
255
256
257
            else:
                logger.debug(
                    "Store the %s-th graph for shape %s from %s via handle %s",
258
259
260
261
262
                    graph_index,
                    str(runtime_shape),
                    self.compiler.name,
                    handle,
                )
263
264
265
266
267
268
269

        # 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:
270
                logger.info("Compiling a graph for dynamic shape takes %.2f s", elapsed)
271
            else:
272
273
274
275
276
                logger.info(
                    "Compiling a graph for shape %s takes %.2f s",
                    runtime_shape,
                    elapsed,
                )
277

278
        return compiled_graph
279
280


281
282
283
@dataclasses.dataclass
class SplitItem:
    submod_name: str
284
    graph_id: int
285
286
287
288
    is_splitting_graph: bool
    graph: fx.GraphModule


289
def split_graph(
290
    graph: fx.GraphModule, resolved_ops: list[torch._ops.OpOverload]
291
) -> tuple[fx.GraphModule, list[SplitItem]]:
292
293
294
295
296
297
298
    # 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
299
300
301
302
303
304
        # Match node.target against resolved_ops
        # node.target can be OpOverloadPacket, need to check .default
        if node.op == "call_function" and (
            node.target in resolved_ops
            or (hasattr(node.target, "default") and node.target.default in resolved_ops)
        ):
305
306
307
308
309
310
311
312
313
314
315
316
            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(
317
318
        graph, None, lambda node: node_to_subgraph_id[node], keep_original_order=True
    )
319

320
    outputs = []
321

322
    names = [name for (name, module) in split_gm.named_modules()]
323

324
325
326
327
    for name in names:
        if "." in name or name == "":
            # recursive child module or the root module
            continue
328

329
        module = getattr(split_gm, name)
330

331
        graph_id = int(name.replace("submod_", ""))
332
        outputs.append(SplitItem(name, graph_id, (graph_id in split_op_graphs), module))
333

334
    # sort by integer graph_id, rather than string name
335
    outputs.sort(key=lambda x: x.graph_id)
336

337
    return split_gm, outputs
338
339


340
341
compilation_start_time = 0.0

342
343
344
345
346
347

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.
348
349
350
351
352

    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.
353
354
    """

355
356
357
358
359
360
361
    def __init__(
        self,
        module: torch.fx.GraphModule,
        compile_submod_names: list[str],
        vllm_config: VllmConfig,
        vllm_backend: "VllmBackend",
    ):
362
363
        super().__init__(module)
        from torch._guards import detect_fake_mode
364

365
366
        self.fake_mode = detect_fake_mode()
        self.compile_submod_names = compile_submod_names
367
368
        self.compilation_config = vllm_config.compilation_config
        self.vllm_config = vllm_config
369
        self.vllm_backend = vllm_backend
370
371
        # When True, it annoyingly dumps the torch.fx.Graph on errors.
        self.extra_traceback = False
372
373
374
375
376
377

    def run(self, *args):
        fake_args = [
            self.fake_mode.from_tensor(t) if isinstance(t, torch.Tensor) else t
            for t in args
        ]
378
        with self.fake_mode, enable_python_dispatcher():
379
            return super().run(*fake_args)
380

381
382
383
384
385
386
    def call_module(
        self,
        target: torch.fx.node.Target,
        args: tuple[torch.fx.node.Argument, ...],
        kwargs: dict[str, Any],
    ) -> Any:
387
388
389
390
        assert isinstance(target, str)
        output = super().call_module(target, args, kwargs)

        if target in self.compile_submod_names:
391
            index = self.compile_submod_names.index(target)
392
393
394
395
            submod = self.fetch_attr(target)
            sym_shape_indices = [
                i for i, x in enumerate(args) if isinstance(x, torch.SymInt)
            ]
396
            global compilation_start_time
397

398
399
400
401
402
403
404
405
406
407
408
            compiled_graph_for_dynamic_shape = (
                self.vllm_backend.compiler_manager.compile(
                    submod,
                    args,
                    self.compilation_config.inductor_compile_config,
                    self.compilation_config,
                    graph_index=index,
                    num_graphs=len(self.compile_submod_names),
                    runtime_shape=None,
                )
            )
409
            # Lazy import here to avoid circular import
410
            from .piecewise_backend import PiecewiseBackend
411

412
            piecewise_backend = PiecewiseBackend(
413
414
415
416
417
418
419
420
                submod,
                self.vllm_config,
                index,
                len(self.compile_submod_names),
                sym_shape_indices,
                compiled_graph_for_dynamic_shape,
                self.vllm_backend,
            )
421

422
423
424
425
            if (
                self.compilation_config.cudagraph_mode.has_piecewise_cudagraphs()
                and not self.compilation_config.use_inductor_graph_partition
            ):
426
427
428
429
                # We're using Dynamo-based piecewise splitting, so we wrap
                # the whole subgraph with a static graph wrapper.
                from .cuda_graph import CUDAGraphOptions

430
431
432
                # resolve the static graph wrapper class (e.g. CUDAGraphWrapper
                # class) as platform dependent.
                static_graph_wrapper_class = resolve_obj_by_qualname(
433
434
                    current_platform.get_static_graph_wrapper_cls()
                )
435
436
437
438
439
440
441
442
443
444
445
446

                # Always assign PIECEWISE runtime mode to the
                # CUDAGraphWrapper for piecewise_backend, to distinguish
                # it from the FULL cudagraph runtime mode, no matter it
                # is wrapped on a full or piecewise fx graph.
                self.module.__dict__[target] = static_graph_wrapper_class(
                    runnable=piecewise_backend,
                    vllm_config=self.vllm_config,
                    runtime_mode=CUDAGraphMode.PIECEWISE,
                    cudagraph_options=CUDAGraphOptions(
                        debug_log_enable=piecewise_backend.is_first_graph,
                        gc_disable=not piecewise_backend.is_first_graph,
447
448
449
                        weak_ref_output=piecewise_backend.is_last_graph,
                    ),
                )
450
451
452
            else:
                self.module.__dict__[target] = piecewise_backend

453
454
455
456
457
            compilation_counter.num_piecewise_capturable_graphs_seen += 1

        return output


458
459
460
461
462
463
464
465
466
# the tag for the part of model being compiled,
# e.g. backbone/eagle_head
model_tag: str = "backbone"


@contextmanager
def set_model_tag(tag: str):
    """Context manager to set the model tag."""
    global model_tag
467
    assert tag != model_tag, (
468
        f"Model tag {tag} is the same as the current tag {model_tag}."
469
    )
470
471
472
473
474
475
476
477
    old_tag = model_tag
    model_tag = tag
    try:
        yield
    finally:
        model_tag = old_tag


478
class VllmBackend:
479
    """The compilation backend for `torch.compile` with vLLM.
480
481
    It is used for compilation level of `CompilationLevel.PIECEWISE`,
    where we customize the compilation.
482

483
484
    The major work of this backend is to split the graph into
    piecewise graphs, and pass them to the piecewise backend.
485

486
487
    This backend also adds the PostGradPassManager to Inductor config,
    which handles the post-grad passes.
488
    """
489

490
491
    vllm_config: VllmConfig
    compilation_config: CompilationConfig
492
493
494
495
496
    _called: bool = False
    # the graph we compiled
    graph: fx.GraphModule
    # the stiching graph module for all the piecewise graphs
    split_gm: fx.GraphModule
497
    piecewise_graphs: list[SplitItem]
498
    returned_callable: Callable
499
500
    # Inductor passes to run on the graph pre-defunctionalization
    post_grad_passes: Sequence[Callable]
501
502
    sym_tensor_indices: list[int]
    input_buffers: list[torch.Tensor]
503
    compiler_manager: CompilerManager
504

505
506
    def __init__(
        self,
507
        vllm_config: VllmConfig,
508
        prefix: str = "",
509
    ):
510
511
        # if the model is initialized with a non-empty prefix,
        # then usually it's enough to use that prefix,
512
        # e.g. language_model, vision_model, etc.
513
514
515
516
517
        # when multiple parts are initialized as independent
        # models, we need to use the model_tag to distinguish
        # them, e.g. backbone (default), eagle_head, etc.
        self.prefix = prefix or model_tag

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

521
522
523
        self.sym_tensor_indices = []
        self.input_buffers = []

524
525
        self.vllm_config = vllm_config
        self.compilation_config = vllm_config.compilation_config
526

527
        self.compiler_manager: CompilerManager = CompilerManager(
528
529
            self.compilation_config
        )
530

531
532
533
        # `torch.compile` is JIT compiled, so we don't need to
        # do anything here

534
    def configure_post_pass(self):
535
        config = self.compilation_config
536
        self.post_grad_pass_manager.configure(self.vllm_config)
537

538
539
        # 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.
540
        inductor_config = config.inductor_compile_config
541
542
        PASS_KEY = "post_grad_custom_post_pass"
        if PASS_KEY in inductor_config:
543
            if isinstance(inductor_config[PASS_KEY], PostGradPassManager):
544
                # PassManager already added to config, make sure it's correct
545
546
547
548
                assert (
                    inductor_config[PASS_KEY].uuid()
                    == self.post_grad_pass_manager.uuid()
                )
549
            else:
550
                # Config should automatically wrap all inductor passes
551
552
                assert isinstance(inductor_config[PASS_KEY], InductorPass)
                self.post_grad_pass_manager.add(inductor_config[PASS_KEY])
553
        inductor_config[PASS_KEY] = self.post_grad_pass_manager
554

555
556
557
558
559
    def __call__(
        self, graph: fx.GraphModule, example_inputs
    ) -> VllmSerializableFunction:
        from .caching import _compute_code_hash, compilation_config_hash_factors

560
        vllm_config = self.vllm_config
561
562
563
564
565
566
        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.

567
            factors = compilation_config_hash_factors(vllm_config)
568
569
            # 2. factors come from the code files that are traced by Dynamo (
            #    it mainly summarizes how the model is used in forward pass)
570
            code_hash = _compute_code_hash(self.compilation_config.traced_files)
571
            self.compilation_config.traced_files.clear()
572
573
574
575
576
577
578
            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
579
580
581
            hash_key = hashlib.md5(
                str(factors).encode(), usedforsecurity=False
            ).hexdigest()[:10]
582
583

            cache_dir = os.path.join(
584
585
586
587
588
589
                envs.VLLM_CACHE_ROOT,
                "torch_compile_cache",
                hash_key,
            )
            self.compilation_config.cache_dir = cache_dir

590
        cache_dir = self.compilation_config.cache_dir
591
        os.makedirs(cache_dir, exist_ok=True)
592
        self.compilation_config.cache_dir = cache_dir
593
594
        rank = vllm_config.parallel_config.rank
        dp_rank = vllm_config.parallel_config.data_parallel_rank
595
        local_cache_dir = os.path.join(cache_dir, f"rank_{rank}_{dp_rank}", self.prefix)
596
        os.makedirs(local_cache_dir, exist_ok=True)
597
        self.compilation_config.local_cache_dir = local_cache_dir
598

599
600
601
        disable_cache = envs.VLLM_DISABLE_COMPILE_CACHE

        if disable_cache:
602
603
            logger.info("vLLM's torch.compile cache is disabled.")
        else:
604
605
606
            logger.info(
                "Using cache directory: %s for vLLM's torch.compile", local_cache_dir
            )
607

608
609
610
        self.compiler_manager.initialize_cache(
            local_cache_dir, disable_cache, self.prefix
        )
611

612
613
        # when dynamo calls the backend, it means the bytecode
        # transform and analysis are done
614
        compilation_counter.num_graphs_seen += 1
615
        from .monitor import torch_compile_start_time
616

617
618
        dynamo_time = time.time() - torch_compile_start_time
        logger.info("Dynamo bytecode transform time: %.2f s", dynamo_time)
619
        self.compilation_config.compilation_time += dynamo_time
620
621
622
623
624
625

        # 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
626
        self.configure_post_pass()
627

628
629
630
631
632
633
634
635
        if self.compilation_config.use_inductor_graph_partition:
            # Let Inductor decide partitioning; avoid FX-level pre-splitting.
            fx_split_ops: list[str] = []
        else:
            fx_split_ops = self.compilation_config.splitting_ops or []

        resolved_split_ops = resolve_defined_ops(fx_split_ops)
        self.split_gm, self.piecewise_graphs = split_graph(graph, resolved_split_ops)
636

637
        from torch._dynamo.utils import lazy_format_graph_code
638
639
640
641
642

        # 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)
643

644
        compilation_counter.num_piecewise_graphs_seen += len(self.piecewise_graphs)
645
        submod_names_to_compile = [
646
647
            item.submod_name
            for item in self.piecewise_graphs
648
649
650
651
652
            if not item.is_splitting_graph
        ]

        # propagate the split graph to the piecewise backend,
        # compile submodules with symbolic shapes
653
654
655
        PiecewiseCompileInterpreter(
            self.split_gm, submod_names_to_compile, self.vllm_config, self
        ).run(*example_inputs)
656

657
658
659
660
        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
661
662
663
664
            src = (
                "from __future__ import annotations\nimport torch\n"
                + self.split_gm.print_readable(print_output=False)
            )
665
666
667
668
669
670
            src = src.replace("<lambda>", "GraphModule")
            with open(graph_path, "w") as f:
                f.write(src)

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

671
672
        self._called = True

673
674
675
676
        if (
            self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
            or not self.compilation_config.cudagraph_copy_inputs
        ):
677
678
679
            return VllmSerializableFunction(
                graph, example_inputs, self.prefix, self.split_gm
            )
680
681
682

        # if we need to copy input buffers for cudagraph
        from torch._guards import detect_fake_mode
683

684
685
686
687
688
689
690
        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)
691
692
693
        # 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
694

695
        self.sym_tensor_indices = [
696
697
698
699
            i
            for i, x in enumerate(fake_args)
            if isinstance(x, torch._subclasses.fake_tensor.FakeTensor)
            and any(is_symbolic(d) for d in x.size())
700
701
702
703
704
705
706
707
708
        ]

        # 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
709
710
        # this is the callable we return to Dynamo to run
        def copy_and_call(*args):
711
712
713
714
715
716
717
718
719
720
721
722
723
            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)

724
725
726
        return VllmSerializableFunction(
            graph, example_inputs, self.prefix, copy_and_call
        )