compiler_interface.py 23.3 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
104
    def load(
        self,
        handle: Any,
        graph: fx.GraphModule,
        example_inputs: list[Any],
        graph_index: int,
105
        compile_range: Range,
106
    ) -> Callable[..., Any]:
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
141
        """
        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:
142
        self.guards: list[Any] = []
143

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

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

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


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

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

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

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


170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
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
    )


187
188
189
190
191
192
193
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.

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

197
198
    name = "inductor_standalone"

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

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

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

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

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

        from torch._inductor import standalone_compile
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
271
        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

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

        if is_compile_cache_enabled(compiler_config):
277
            compiled_graph.save(path=path, format=self.save_format)
278
            compilation_counter.num_compiled_artifacts_saved += 1
279
280
        return compiled_graph, (key, path)

281
282
283
284
285
286
    def load(
        self,
        handle: Any,
        graph: fx.GraphModule,
        example_inputs: list[Any],
        graph_index: int,
287
        compile_range: Range,
288
    ) -> Callable[..., Any]:
289
290
291
292
293
        assert isinstance(handle, tuple)
        assert isinstance(handle[0], str)
        assert isinstance(handle[1], str)
        path = handle[1]
        inductor_compiled_graph = torch._inductor.CompiledArtifact.load(
294
            path=path, format=self.save_format
295
        )
296
        from torch._inductor.compile_fx import graph_returns_tuple
297

298
299
        returns_tuple = graph_returns_tuple(graph)

300
        def compiled_graph_wrapper(*args: Any) -> tuple[Any, ...] | Any:
301
302
303
304
305
306
307
308
309
310
311
312
            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


313
314
class InductorAdaptor(CompilerInterface):
    """
315
    The adaptor for the Inductor compiler, version 2.5, 2.6, 2.7.
316
    """
317

318
319
320
    name = "inductor"

    def compute_hash(self, vllm_config: VllmConfig) -> str:
321
        factors = get_inductor_factors()
322
323
324
        hash_str: str = safe_hash(
            str(factors).encode(), usedforsecurity=False
        ).hexdigest()[:10]
325
326
        return hash_str

327
328
    def initialize_cache(
        self, cache_dir: str, disable_cache: bool = False, prefix: str = ""
329
    ) -> None:
330
        self.cache_dir = cache_dir
331
        self.prefix = prefix
332
        self.base_cache_dir = cache_dir[: -len(prefix)] if prefix else cache_dir
333
334
        if disable_cache:
            return
335
        # redirect the cache directory to a subdirectory
336
337
338
        # 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.
339
        inductor_cache = os.path.join(self.base_cache_dir, "inductor_cache")
340
341
        os.makedirs(inductor_cache, exist_ok=True)
        os.environ["TORCHINDUCTOR_CACHE_DIR"] = inductor_cache
342
        triton_cache = os.path.join(self.base_cache_dir, "triton_cache")
343
344
345
346
347
348
        os.makedirs(triton_cache, exist_ok=True)
        os.environ["TRITON_CACHE_DIR"] = triton_cache

    def compile(
        self,
        graph: fx.GraphModule,
349
350
        example_inputs: list[Any],
        compiler_config: dict[str, Any],
351
        compile_range: Range,
352
        key: str | None = None,
353
    ) -> tuple[Callable[..., Any] | None, Any | None]:
354
        compilation_counter.num_inductor_compiles += 1
355
        from torch._inductor.compile_fx import compile_fx
356

357
358
359
        current_config = {}
        if compiler_config is not None:
            current_config.update(compiler_config)
360
361
362
363
364

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

365
        set_inductor_config(current_config, compile_range)
366
        set_functorch_config()
367
368
369
370
371
372
373
374
375
376
377

        # 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
378
        from torch._inductor.codecache import compiled_fx_graph_hash
379

380
381
382
383
384
        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:
385
                nonlocal file_path
386
387
                compiled_fn = inductor_compiled_graph.current_callable
                file_path = compiled_fn.__code__.co_filename  # noqa
388
389
390
391
                if (
                    not file_path.startswith(self.base_cache_dir)
                    and compiled_fn.__closure__ is not None
                ):
392
393
394
395
396
                    # 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
397
398
399
400
401
                        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
402
                            break
403
404
                hash_str = inductor_compiled_graph._fx_graph_cache_key
            return output
405

406
        def hijack_compiled_fx_graph_hash(*args: Any, **kwargs: Any) -> Any:
407
408
409
410
411
            out = compiled_fx_graph_hash(*args, **kwargs)
            nonlocal hash_str
            hash_str = out[0]
            return out

412
        def _check_can_cache(*args: Any, **kwargs: Any) -> None:
413
414
415
416
417
418
419
420
421
422
423
424
425
426
            # 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(
427
428
429
430
431
                patch(
                    "torch._inductor.codecache.compiled_fx_graph_hash",
                    hijack_compiled_fx_graph_hash,
                )
            )
432
433
434

            # for providing a dummy shape environment
            stack.enter_context(
435
436
437
438
439
                patch(
                    "torch._inductor.codecache.FxGraphCache._get_shape_env",
                    _get_shape_env,
                )
            )
440

441
            from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
442
443
444
445
446
447

            # 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",
448
449
450
                        _get_shape_env,
                    )
                )
451

452
453
454
455
            # for forcing the graph to be cached
            stack.enter_context(
                patch(
                    "torch._inductor.codecache.FxGraphCache._check_can_cache",
456
457
458
                    _check_can_cache,
                )
            )
459

460
461
462
            # Dynamo metrics context, see method for more details.
            stack.enter_context(self.metrics_context())

463
464
465
466
467
468
            # 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.
469
470
471
472
473
474
475
476
477
478
479
480
481
            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)
            )
482

483
484
485
486
487
488
            compiled_graph = compile_fx(
                graph,
                example_inputs,
                inner_compile=hijacked_compile_fx_inner,
                config_patches=current_config,
            )
489

490
491
        # Turn off the checks if we disable the compilation cache.
        if is_compile_cache_enabled(compiler_config):
492
493
494
495
496
497
498
            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 "
499
500
                    "to see the real issue. "
                )
501
            assert file_path is not None, (
502
503
                "failed to get the file path of the compiled graph"
            )
504
505
        return compiled_graph, (hash_str, file_path)

506
507
508
509
510
511
    def load(
        self,
        handle: Any,
        graph: fx.GraphModule,
        example_inputs: list[Any],
        graph_index: int,
512
        compile_range: Range,
513
    ) -> Callable[..., Any]:
514
515
516
517
518
        assert isinstance(handle, tuple)
        assert isinstance(handle[0], str)
        assert isinstance(handle[1], str)
        hash_str = handle[0]

519
        from torch._functorch._aot_autograd.autograd_cache import AOTAutogradCache
520
        from torch._inductor.codecache import FxGraphCache
521

522
523
        with ExitStack() as exit_stack:
            exit_stack.enter_context(
524
525
526
527
528
                patch(
                    "torch._inductor.codecache.FxGraphCache._get_shape_env",
                    lambda *args, **kwargs: AlwaysHitShapeEnv(),
                )
            )
529
530
531
532
533
            # 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",
534
535
536
                        lambda *args, **kwargs: AlwaysHitShapeEnv(),
                    )
                )
537
538
539
540

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

541
            from torch._inductor.output_code import CompiledFxGraphConstantsWithGm
542

543
544
545
546
547
548
549
550
            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
            )
551
552
553
554
555
556
557
558

        # 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
559

560
561
562
        returns_tuple = graph_returns_tuple(graph)

        # this is the callable we return to Dynamo to run
563
        def compiled_graph(*args: Any) -> tuple[Any, ...] | Any:
564
565
566
567
568
569
570
571
572
573
574
            # 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

575
    def metrics_context(self) -> contextlib.AbstractContextManager[Any]:
576
577
578
579
580
581
582
583
584
        """
        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
585
        should be set at a different mode of compilation.
586
587
588
589
590

        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.
        """
591
        if is_torch_equal_or_newer("2.6"):
592
            import torch._dynamo.utils
593

594
            return torch._dynamo.utils.get_metrics_context()  # type: ignore[no-any-return]
595
596
597
        else:
            return contextlib.nullcontext()

598

599
def set_inductor_config(config: dict[str, Any], compile_range: Range) -> None:
600
601
    if compile_range.is_single_size():
        # for a specific batch size, tuning triton kernel parameters
602
        # can be beneficial
603
604
        config["max_autotune"] = envs.VLLM_ENABLE_INDUCTOR_MAX_AUTOTUNE
        config["coordinate_descent_tuning"] = (
605
606
            envs.VLLM_ENABLE_INDUCTOR_COORDINATE_DESCENT_TUNING
        )
607
608


609
def set_functorch_config() -> None:
610
611
    if not envs.VLLM_USE_MEGA_AOT_ARTIFACT:
        torch._functorch.config.bundled_autograd_cache = False
612
613


614
615
616
617
618
619
class EagerAdaptor(CompilerInterface):
    name = "eager"

    def compile(
        self,
        graph: fx.GraphModule,
620
621
        example_inputs: list[Any],
        compiler_config: dict[str, Any],
622
        compile_range: Range,
623
        key: str | None = None,
624
    ) -> tuple[Callable[..., Any] | None, Any | None]:
625
        compilation_counter.num_eager_compiles += 1
626
627
628
        # 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