pass_manager.py 2.69 KB
Newer Older
1
2
# SPDX-License-Identifier: Apache-2.0

3
4
from torch import fx as fx

5
from vllm.config import VllmConfig
6
7
from vllm.logger import init_logger

8
from .activation_quant_fusion import ActivationQuantFusionPass
9
10
from .fix_functionalization import FixFunctionalizationPass
from .fusion import FusionPass
11
from .inductor_pass import CustomGraphPass, InductorPass, get_pass_context
12
from .noop_elimination import NoOpEliminationPass
13
14
from .sequence_parallelism import SequenceParallelismPass
from .vllm_inductor_pass import VllmInductorPass
15
16
17
18

logger = init_logger(__name__)


19
class PostGradPassManager(CustomGraphPass):
20
21
22
    """
    The pass manager for post-grad passes.
    It handles configuration, adding custom passes, and running passes.
23
24
    It supports uuid for the Inductor code cache. That includes torch<2.6
    support using pickling (in .inductor_pass.CustomGraphPass).
25
26
27

    The order of the post-grad post-passes is:
    1. passes (constructor parameter)
28
    2. default passes (NoopEliminationPass, FusionPass)
29
30
31
32
33
34
    3. config["post_grad_custom_post_pass"] (if it exists)
    4. fix_functionalization
    This way, all passes operate on a functionalized graph.
    """

    def __init__(self):
35
        self.passes: list[VllmInductorPass] = []
36
37

    def __call__(self, graph: fx.Graph):
38
        shape = get_pass_context().runtime_shape
39
        for pass_ in self.passes:
40
41
            if pass_.is_applicable_for_shape(shape):
                pass_(graph)
42
43
44
45

        # always run fix_functionalization last
        self.fix_functionalization(graph)

46
47
48
49
    def configure(self, config: VllmConfig):
        self.pass_config = config.compilation_config.pass_config
        if self.pass_config.enable_noop:
            self.passes += [NoOpEliminationPass(config)]
50

51
52
        if self.pass_config.enable_fusion:
            self.passes += [FusionPass.instance(config)]
53
            self.passes += [ActivationQuantFusionPass(config)]
54

55
56
57
58
        if self.pass_config.enable_sequence_parallelism:
            self.passes += [SequenceParallelismPass(config)]

        self.fix_functionalization = FixFunctionalizationPass(config)
59
60
61
62
63

    def add(self, pass_: InductorPass):
        assert isinstance(pass_, InductorPass)
        self.passes.append(pass_)

64
    def uuid(self):
65
        """
66
67
68
        The PostGradPassManager is set as a custom pass in the Inductor and
        affects compilation caching. Its uuid depends on the UUIDs of all
        dependent passes and the pass config. See InductorPass for more info.
69
70
71
72
73
        """
        state = {"pass_config": self.pass_config.uuid(), "passes": []}
        for pass_ in self.passes:
            state["passes"].append(pass_.uuid())
        state["passes"].append(self.fix_functionalization.uuid())
74
        return InductorPass.hash_dict(state)