compiler_interface.py 24.2 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
import contextlib
4
5
import copy
import os
6
from collections.abc import Callable
7
from contextlib import ExitStack
8
from typing import Any, Literal
9
10
11
12
13
14
from unittest.mock import patch

import torch
import torch._inductor.compile_fx
import torch.fx as fx

15
import vllm.envs as envs
16
from vllm.compilation.counter import compilation_counter
17
from vllm.config import VllmConfig
18
from vllm.config.utils import Range
19
from vllm.logger import init_logger
20
from vllm.utils.hashing import safe_hash
21
from vllm.utils.torch_utils import is_torch_equal_or_newer
22

23
24
logger = init_logger(__name__)

25
26
27
28
29

class CompilerInterface:
    """
    The interface for a compiler that can be used by vLLM.
    """
30

31
32
33
34
    # The name of the compiler, e.g. inductor.
    # This is a class-level attribute.
    name: str

35
36
    def initialize_cache(
        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
37
    ) -> None:
38
39
40
41
        """
        when the vLLM process uses `cache_dir` as the cache directory,
        the compiler should initialize itself with the cache directory,
        e.g. by re-directing its own cache directory to a sub-directory.
42
43
44
45
46
47
48
49

        prefix can be used in combination with cache_dir to figure out the base
        cache directory, e.g. there're multiple parts of model being compiled,
        but we want to share the same cache directory for all of them.

        e.g.
        cache_dir = "/path/to/dir/backbone", prefix = "backbone"
        cache_dir = "/path/to/dir/eagle_head", prefix = "eagle_head"
50
51
52
53
54
        """
        pass

    def compute_hash(self, vllm_config: VllmConfig) -> str:
        """
55
        Gather all the relevant information from the vLLM config,
56
57
        to compute a hash so that we can cache the compiled model.

58
59
        See [`VllmConfig.compute_hash`][vllm.config.VllmConfig.compute_hash]
        to check what information
60
61
62
63
64
65
66
67
        is already considered by default. This function should only
        consider the information that is specific to the compiler.
        """
        return ""

    def compile(
        self,
        graph: fx.GraphModule,
68
69
        example_inputs: list[Any],
        compiler_config: dict[str, Any],
70
        compile_range: Range,
71
        key: str | None = None,
72
    ) -> tuple[Callable[..., Any] | None, Any | None]:
73
74
        """
        Compile the graph with the given example inputs and compiler config,
75
76
77
78
79
        with a range. The `compile_range` specifies the range of the inputs,
        it could be concrete size (if compile_sizes is provided), e.g. [4, 4]
        or a range [5, 8].
        Right now we only support one variable in ranges for all inputs,
         which is the batchsize (number of tokens) during inference.
80
81
82
83
84
85
86
87
88
89
90
91

        Dynamo will make sure `graph(*example_inputs)` is valid.

        The function should return a compiled callable function, as well as
        a handle that can be used to directly load the compiled function.

        The handle should be a plain Python object, preferably a string or a
        file path for readability.

        If the compiler doesn't support caching, it should return None for the
        handle. If the compiler fails to compile the graph, it should return
        None for the compiled function as well.
92
93
94
95

        `key` is required for StandaloneInductorAdapter, it specifies where to
        save the compiled artifact. The compiled artifact gets saved to
        `cache_dir/key`.
96
97
98
        """
        return None, None

99
100
101
102
103
    def load(
        self,
        handle: Any,
        graph: fx.GraphModule,
        example_inputs: list[Any],
104
        compile_range: Range,
105
    ) -> Callable[..., Any]:
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
        """
        Load the compiled function from the handle.
        Raises an error if the handle is invalid.

        The handle is the second return value of the `compile` function.
        """
        raise NotImplementedError("caching is not supported")


class AlwaysHitShapeEnv:
    """
    Why do we need this class:

    For normal `torch.compile` usage, every compilation will have
    one Dynamo bytecode compilation and one Inductor compilation.
    The Inductor compilation happens under the context of the
    Dynamo bytecode compilation, and that context is used to
    determine the dynamic shape information, etc.

    For our use case, we only run Dynamo bytecode compilation once,
    and run Inductor compilation multiple times with different shapes
    plus a general shape. The compilation for specific shapes happens
    outside of the context of the Dynamo bytecode compilation. At that
    time, we don't have shape environment to provide to Inductor, and
    it will fail the Inductor code cache lookup.

    By providing a dummy shape environment that always hits, we can
    make the Inductor code cache lookup always hit, and we can
    compile the graph for different shapes as needed.

    The following dummy methods are obtained by trial-and-error
    until it works.
    """

    def __init__(self) -> None:
141
        self.guards: list[Any] = []
142

143
    def evaluate_guards_expression(self, *args: Any, **kwargs: Any) -> Literal[True]:
144
145
        return True

146
    def get_pruned_guards(self, *args: Any, **kwargs: Any) -> list[Any]:
147
148
        return []

149
    def produce_guards_expression(self, *args: Any, **kwargs: Any) -> Literal[""]:
150
151
152
        return ""


153
154
def get_inductor_factors() -> list[Any]:
    factors: list[Any] = []
155
156
    # summarize system state
    from torch._inductor.codecache import CacheBase
157

158
159
160
161
162
    system_factors = CacheBase.get_system()
    factors.append(system_factors)

    # summarize pytorch state
    from torch._inductor.codecache import torch_key
163

164
165
166
167
168
    torch_factors = torch_key()
    factors.append(torch_factors)
    return factors


169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
def is_compile_cache_enabled(
    vllm_additional_inductor_config: dict[str, Any],
) -> bool:
    vllm_inductor_config_disable_cache = vllm_additional_inductor_config.get(
        "force_disable_caches", False
    )

    # TODO(gmagogsfm): Replace torch._inductor.config.force_disable_caches
    # with torch.compiler.config.force_disable_caches when minimum PyTorch
    # version reaches 2.10
    return (
        not envs.VLLM_DISABLE_COMPILE_CACHE
        and not torch._inductor.config.force_disable_caches
        and not vllm_inductor_config_disable_cache
    )


186
187
188
189
190
191
192
class InductorStandaloneAdaptor(CompilerInterface):
    """
    The adaptor for the Inductor compiler.
    Requires PyTorch 2.8+.
    This is not on by default yet, but we plan to turn it on by default for
    PyTorch 2.8.

193
    Use VLLM_USE_STANDALONE_COMPILE to toggle this on or off.
194
    """
195

196
197
    name = "inductor_standalone"

198
    def __init__(self, save_format: Literal["binary", "unpacked"]) -> None:
199
200
        self.save_format = save_format

201
202
    def compute_hash(self, vllm_config: VllmConfig) -> str:
        factors = get_inductor_factors()
203
204
205
        hash_str: str = safe_hash(
            str(factors).encode(), usedforsecurity=False
        ).hexdigest()[:10]
206
207
        return hash_str

208
209
    def initialize_cache(
        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
210
    ) -> None:
211
212
213
214
215
        self.cache_dir = cache_dir

    def compile(
        self,
        graph: fx.GraphModule,
216
217
        example_inputs: list[Any],
        compiler_config: dict[str, Any],
218
        compile_range: Range,
219
        key: str | None = None,
220
    ) -> tuple[Callable[..., Any] | None, Any | None]:
221
        compilation_counter.num_inductor_compiles += 1
222
223
224
        current_config = {}
        if compiler_config is not None:
            current_config.update(compiler_config)
225
        set_inductor_config(current_config, compile_range)
226
        set_functorch_config()
227

228
        if compile_range.is_single_size():
229
230
            dynamic_shapes = "from_example_inputs"
        else:
231
            dynamic_shapes = "from_graph"
232
233

        from torch._inductor import standalone_compile
234

235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
        supports_aot = is_torch_equal_or_newer("2.10.0.dev")

        if not supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT:
            logger.error(
                "CRITICAL: VLLM_USE_MEGA_AOT_ARTIFACT "
                "is enabled but PyTorch version does not support 'aot' "
                "parameter in standalone_compile. This requires PyTorch "
                "2.10.0+. Falling back to non-AOT mode."
            )

        compile_kwargs = {
            "dynamic_shapes": dynamic_shapes,
            "options": {
                "config_patches": current_config,
            },
        }

        use_aot: bool = supports_aot and envs.VLLM_USE_MEGA_AOT_ARTIFACT
        # only add 'aot' parameter if both supported and enabled...
        # this will set bundled_autograd_cache
        # https://github.com/pytorch/pytorch/blob/9bbc5b2905c260adf41bc866a732f9c121a2828a/torch/_inductor/standalone_compile.py#L359 # noqa
        if use_aot:
            compile_kwargs["aot"] = True  # type: ignore[assignment]

        compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)

        if use_aot:
            from torch._inductor.standalone_compile import AOTCompiledArtifact

            assert isinstance(compiled_graph, AOTCompiledArtifact)
            assert hasattr(compiled_graph, "serialize")
            # just return the compiled graph and a key
            # since we can serialize the bytes using to_bytes
            # and reload it using the key when reading
            return compiled_graph, None

271
272
273
        # Save the compiled artifact to disk in the specified path
        assert key is not None
        path = os.path.join(self.cache_dir, key)
274

275
276
277
278
279
280
281
        def is_saveable_2_10(compiled_artifact):
            # can just use compiled_artifact.is_saveable in 2.11
            if compiled_artifact._artifacts is None:
                return False
            _, cache_info = compiled_artifact._artifacts
            return len(cache_info.aot_autograd_artifacts) == 1

282
        if is_compile_cache_enabled(compiler_config):
283
284
285
286
287
288
289
290
291
292
293
294
            if not is_saveable_2_10(compiled_graph):
                raise RuntimeError(
                    "The compiled artifact is not serializable. This usually means "
                    "that the model code has something that is not serializable "
                    "by torch.compile in it. You can fix this by either "
                    "figuring out what is not serializable and rewriting it, "
                    "filing a bug report, "
                    "or suppressing this error by "
                    "disabling vLLM's compilation cache via "
                    "VLLM_DISABLE_COMPILE_CACHE=1 "
                    "(this will greatly increase vLLM server warm start times)."
                )
295
            compiled_graph.save(path=path, format=self.save_format)
296
            compilation_counter.num_compiled_artifacts_saved += 1
297
298
        return compiled_graph, (key, path)

299
300
301
302
303
    def load(
        self,
        handle: Any,
        graph: fx.GraphModule,
        example_inputs: list[Any],
304
        compile_range: Range,
305
    ) -> Callable[..., Any]:
306
307
308
309
310
        assert isinstance(handle, tuple)
        assert isinstance(handle[0], str)
        assert isinstance(handle[1], str)
        path = handle[1]
        inductor_compiled_graph = torch._inductor.CompiledArtifact.load(
311
            path=path, format=self.save_format
312
        )
313
        from torch._inductor.compile_fx import graph_returns_tuple
314

315
316
        returns_tuple = graph_returns_tuple(graph)

317
        def compiled_graph_wrapper(*args: Any) -> tuple[Any, ...] | Any:
318
319
320
321
322
323
324
325
326
327
328
329
            graph_output = inductor_compiled_graph(*args)
            # unpack the tuple if needed
            # TODO(rzou): the implication is that we're not
            # reading the python bytecode correctly in vLLM?
            if returns_tuple:
                return graph_output
            else:
                return graph_output[0]

        return compiled_graph_wrapper


330
331
class InductorAdaptor(CompilerInterface):
    """
332
    The adaptor for the Inductor compiler, version 2.5, 2.6, 2.7.
333
    """
334

335
336
337
    name = "inductor"

    def compute_hash(self, vllm_config: VllmConfig) -> str:
338
        factors = get_inductor_factors()
339
340
341
        hash_str: str = safe_hash(
            str(factors).encode(), usedforsecurity=False
        ).hexdigest()[:10]
342
343
        return hash_str

344
345
    def initialize_cache(
        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
346
    ) -> None:
347
        self.cache_dir = cache_dir
348
        self.prefix = prefix
349
        self.base_cache_dir = cache_dir[: -len(prefix)] if prefix else cache_dir
350
351
        if disable_cache:
            return
352
        # redirect the cache directory to a subdirectory
353
354
355
        # set flags so that Inductor and Triton store their cache
        # in the cache_dir, then users only need to copy the cache_dir
        # to another machine to reuse the cache.
356
        inductor_cache = os.path.join(self.base_cache_dir, "inductor_cache")
357
358
        os.makedirs(inductor_cache, exist_ok=True)
        os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache
359
        triton_cache = os.path.join(self.base_cache_dir, "triton_cache")
360
361
362
363
364
365
        os.makedirs(triton_cache, exist_ok=True)
        os.environ["TRITON_CACHE_DIR"] = triton_cache

    def compile(
        self,
        graph: fx.GraphModule,
366
367
        example_inputs: list[Any],
        compiler_config: dict[str, Any],
368
        compile_range: Range,
369
        key: str | None = None,
370
    ) -> tuple[Callable[..., Any] | None, Any | None]:
371
        compilation_counter.num_inductor_compiles += 1
372
        from torch._inductor.compile_fx import compile_fx
373

374
375
376
        current_config = {}
        if compiler_config is not None:
            current_config.update(compiler_config)
377
378
379
380
381

        # disable remote cache
        current_config["fx_graph_cache"] = True
        current_config["fx_graph_remote_cache"] = False

382
        set_inductor_config(current_config, compile_range)
383
        set_functorch_config()
384
385
386
387
388
389
390
391
392
393
394

        # inductor can inplace modify the graph, so we need to copy it
        # see https://github.com/pytorch/pytorch/issues/138980
        graph = copy.deepcopy(graph)

        # it's the first time we compile this graph
        # the assumption is that we don't have nested Inductor compilation.
        # compiled_fx_graph_hash will only be called once, and we can hook
        # it to get the hash of the compiled graph directly.

        hash_str, file_path = None, None
395
        from torch._inductor.codecache import compiled_fx_graph_hash
396

397
398
399
400
401
        def hijacked_compile_fx_inner(*args: Any, **kwargs: Any) -> Any:
            output = torch._inductor.compile_fx.compile_fx_inner(*args, **kwargs)
            nonlocal hash_str
            inductor_compiled_graph = output
            if inductor_compiled_graph is not None:
402
                nonlocal file_path
403
404
                compiled_fn = inductor_compiled_graph.current_callable
                file_path = compiled_fn.__code__.co_filename  # noqa
405
406
407
408
                if (
                    not file_path.startswith(self.base_cache_dir)
                    and compiled_fn.__closure__ is not None
                ):
409
410
411
412
413
                    # hooked in the align_inputs_from_check_idxs function
                    # in torch/_inductor/utils.py
                    for cell in compiled_fn.__closure__:
                        if not callable(cell.cell_contents):
                            continue
414
415
416
417
418
                        code = cell.cell_contents.__code__
                        if code.co_filename.startswith(self.base_cache_dir):
                            # this is the real file path
                            # compiled from Inductor
                            file_path = code.co_filename
419
                            break
420
421
                hash_str = inductor_compiled_graph._fx_graph_cache_key
            return output
422

423
        def hijack_compiled_fx_graph_hash(*args: Any, **kwargs: Any) -> Any:
424
425
426
427
428
            out = compiled_fx_graph_hash(*args, **kwargs)
            nonlocal hash_str
            hash_str = out[0]
            return out

429
        def _check_can_cache(*args: Any, **kwargs: Any) -> None:
430
431
432
433
434
435
436
437
438
439
440
441
442
443
            # no error means it can be cached.
            # Inductor refuses to cache the graph outside of Dynamo
            # tracing context, and also disables caching for graphs
            # with high-order ops.
            # For vLLM, in either case, we want to cache the graph.
            # see https://github.com/pytorch/pytorch/blob/9f5ebf3fc609105a74eab4ccc24932d6353ff566/torch/_inductor/codecache.py#L1221 # noqa
            return

        def _get_shape_env() -> AlwaysHitShapeEnv:
            return AlwaysHitShapeEnv()

        with ExitStack() as stack:
            # for hijacking the hash of the compiled graph
            stack.enter_context(
444
445
446
447
448
                patch(
                    "torch._inductor.codecache.compiled_fx_graph_hash",
                    hijack_compiled_fx_graph_hash,
                )
            )
449
450
451

            # for providing a dummy shape environment
            stack.enter_context(
452
453
454
455
456
                patch(
                    "torch._inductor.codecache.FxGraphCache._get_shape_env",
                    _get_shape_env,
                )
            )
457

458
            from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
459
460
461
462
463
464

            # torch 2.8+ on main uses _get_shape_env in AOTAutogradCache
            if hasattr(AOTAutogradCache, "_get_shape_env"):
                stack.enter_context(
                    patch(
                        "torch._functorch._aot_autograd.autograd_cache.AOTAutogradCache._get_shape_env",
465
466
467
                        _get_shape_env,
                    )
                )
468

469
470
471
472
            # for forcing the graph to be cached
            stack.enter_context(
                patch(
                    "torch._inductor.codecache.FxGraphCache._check_can_cache",
473
474
475
                    _check_can_cache,
                )
            )
476

477
478
479
            # Dynamo metrics context, see method for more details.
            stack.enter_context(self.metrics_context())

480
481
482
483
484
485
            # Disable remote caching. When these are on, on remote cache-hit,
            # the monkey-patched functions never actually get called.
            # vLLM today assumes and requires the monkey-patched functions to
            # get hit.
            # TODO(zou3519): we're going to replace this all with
            # standalone_compile sometime.
486
487
488
489
490
491
492
493
494
495
496
497
498
            stack.enter_context(
                torch._inductor.config.patch(fx_graph_remote_cache=False)
            )
            # InductorAdaptor (unfortunately) requires AOTAutogradCache
            # to be turned off to run. It will fail to acquire the hash_str
            # and error if not.
            # StandaloneInductorAdaptor (PyTorch 2.8+) fixes this problem.
            stack.enter_context(
                torch._functorch.config.patch(enable_autograd_cache=False)
            )
            stack.enter_context(
                torch._functorch.config.patch(enable_remote_autograd_cache=False)
            )
499

500
501
502
503
504
505
            compiled_graph = compile_fx(
                graph,
                example_inputs,
                inner_compile=hijacked_compile_fx_inner,
                config_patches=current_config,
            )
506

507
508
        # Turn off the checks if we disable the compilation cache.
        if is_compile_cache_enabled(compiler_config):
509
510
511
512
513
514
515
            if hash_str is None:
                raise RuntimeError(
                    "vLLM failed to compile the model. The most "
                    "likely reason for this is that a previous compilation "
                    "failed, leading to a corrupted compilation artifact. "
                    "We recommend trying to "
                    "remove ~/.cache/vllm/torch_compile_cache and try again "
516
517
                    "to see the real issue. "
                )
518
            assert file_path is not None, (
519
520
                "failed to get the file path of the compiled graph"
            )
521
522
        return compiled_graph, (hash_str, file_path)

523
524
525
526
527
    def load(
        self,
        handle: Any,
        graph: fx.GraphModule,
        example_inputs: list[Any],
528
        compile_range: Range,
529
    ) -> Callable[..., Any]:
530
531
532
533
534
        assert isinstance(handle, tuple)
        assert isinstance(handle[0], str)
        assert isinstance(handle[1], str)
        hash_str = handle[0]

535
        from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
536
        from torch._inductor.codecache import FxGraphCache
537

538
539
        with ExitStack() as exit_stack:
            exit_stack.enter_context(
540
541
542
543
544
                patch(
                    "torch._inductor.codecache.FxGraphCache._get_shape_env",
                    lambda *args, **kwargs: AlwaysHitShapeEnv(),
                )
            )
545
546
547
548
549
            # torch 2.8+ on main uses _get_shape_env in AOTAutogradCache
            if hasattr(AOTAutogradCache, "_get_shape_env"):
                exit_stack.enter_context(
                    patch(
                        "torch._functorch._aot_autograd.autograd_cache.AOTAutogradCache._get_shape_env",
550
551
552
                        lambda *args, **kwargs: AlwaysHitShapeEnv(),
                    )
                )
553
554
555
556

            # Dynamo metrics context, see method for more details.
            exit_stack.enter_context(self.metrics_context())

557
            from torch._inductor.output_code import CompiledFxGraphConstantsWithGm
558

559
560
561
562
563
564
565
566
            constants = CompiledFxGraphConstantsWithGm(graph)
            inductor_compiled_graph, _ = FxGraphCache._lookup_graph(
                hash_str, example_inputs, True, None, constants
            )
            assert inductor_compiled_graph is not None, (
                "Inductor cache lookup failed. Please remove "
                f"the cache directory and try again."  # noqa
            )
567
568
569
570
571
572
573
574

        # Inductor calling convention (function signature):
        # f(list) -> tuple
        # Dynamo calling convention (function signature):
        # f(*args) -> Any

        # need to know if the graph returns a tuple
        from torch._inductor.compile_fx import graph_returns_tuple
575

576
577
578
        returns_tuple = graph_returns_tuple(graph)

        # this is the callable we return to Dynamo to run
579
        def compiled_graph(*args: Any) -> tuple[Any, ...] | Any:
580
581
582
583
584
585
586
587
588
589
590
            # convert args to list
            list_args = list(args)
            graph_output = inductor_compiled_graph(list_args)
            # unpack the tuple if needed
            if returns_tuple:
                return graph_output
            else:
                return graph_output[0]

        return compiled_graph

591
    def metrics_context(self) -> contextlib.AbstractContextManager[Any]:
592
593
594
595
596
597
598
599
600
        """
        This method returns the Dynamo metrics context (if it exists,
        otherwise a null context). It is used by various compile components.
        Present in torch>=2.6, it's used inside FxGraphCache in
        torch==2.6 (but not after). It might also be used in various other
        torch.compile internal functions.

        Because it is re-entrant, we always set it (even if entering via Dynamo
        and the context was already entered). We might want to revisit if it
601
        should be set at a different mode of compilation.
602
603
604
605
606

        This is likely a bug in PyTorch: public APIs should not rely on
        manually setting up internal contexts. But we also rely on non-public
        APIs which might not provide these guarantees.
        """
607
        if is_torch_equal_or_newer("2.6"):
608
            import torch._dynamo.utils
609

610
            return torch._dynamo.utils.get_metrics_context()  # type: ignore[no-any-return]
611
612
613
        else:
            return contextlib.nullcontext()

614

615
def set_inductor_config(config: dict[str, Any], compile_range: Range) -> None:
616
617
    if compile_range.is_single_size():
        # for a specific batch size, tuning triton kernel parameters
618
        # can be beneficial
619
620
        config["max_autotune"] = envs.VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE
        config["coordinate_descent_tuning"] = (
621
622
            envs.VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING
        )
623
624


625
def set_functorch_config() -> None:
626
627
    if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
        torch._functorch.config.bundled_autograd_cache = False
628
629


630
631
632
633
634
635
class EagerAdaptor(CompilerInterface):
    name = "eager"

    def compile(
        self,
        graph: fx.GraphModule,
636
637
        example_inputs: list[Any],
        compiler_config: dict[str, Any],
638
        compile_range: Range,
639
        key: str | None = None,
640
    ) -> tuple[Callable[..., Any] | None, Any | None]:
641
        compilation_counter.num_eager_compiles += 1
642
643
644
        # we don't need to compile the graph, just return the graph itself.
        # It does not support caching, return None for the handle.
        return graph, None