Skip to content

vllm.compilation.multi_output_match

MultiOutputMatch

Bases: ABC

This class provides utilities to process multi-output matches and manually insert replacements.

This is necessary because the automatic replacement for multi-output matches is broken: https://github.com/pytorch/pytorch/issues/137280

Source code in vllm/compilation/multi_output_match.py 
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
class MultiOutputMatch(abc.ABC):
    """
    This class provides utilities to process multi-output matches and
    manually insert replacements.

    This is necessary because the automatic replacement for multi-output
    matches is broken: https://github.com/pytorch/pytorch/issues/137280
    """

    def __init__(self, match: pm.Match):
        self.match = match

    @abstractmethod
    def process(self):
        """
        Process a multi-output match and manually insert the replacement.

        This method should:
        1. Insert the replacement nodes after the last node in the match.
        2. Rebind the users of nodes in the match to use the new nodes.
        3. Set meta["val"] for de-functionalization.

        The result of an auto-functionalized node is a tuple of tensors.
        The first element is the return value of the function, usually None.
        The remaining elements are the mutated args of the function.

        All auto-functionalized nodes must contain a proper meta["val"],
        as it is used by de-functionalization. meta["val"] has to contain the
        value of the node (tuple of tensors) that would be returned by the
        functionalized node during tracing.

        Existing nodes in the graph all have this property set, but we have
        to set it manually for new nodes we insert.

        Example:
        # op schema: foo(a: Tensor!, b: Tensor, c: Tensor!) -> None
        at = auto_functionalized(torch.ops._C.foo.default, a, b, c)
        # at.meta["val"] = (None, a, c)
        """
        raise NotImplementedError

    @property
    def nodes(self) -> list[fx.Node]:
        return self.match.nodes

    @property
    def graph(self) -> fx.Graph:
        return self.match.graph

    def find_auto_fn(self, op) -> fx.Node:
        """
        Find the first auto_functionalized node with the given op in the match.
        """
        return find_auto_fn(self.nodes, op)

    def inserting_after_match(self):
        """
        Insert nodes after the last node in the match.
        This is done to avoid use-before-definition errors after inserting
        replacement nodes.
        """

        # match.nodes is not guaranteed to be sorted.
        # Find the last node in the match.
        for last_node_in_match in reversed(self.graph.nodes):
            if last_node_in_match in self.match.nodes:
                break
        else:
            raise ValueError("No nodes in graph")

        return self.graph.inserting_after(last_node_in_match)

    def insert_getitems(self, tuple_node: fx.Node,
                        indices: Iterable[int]) -> tuple[fx.Node, ...]:
        """
        Insert operator.getitem nodes to extract elements from a tuple node.

        :param tuple_node: The tuple node to extract elements from.
        :param indices: The indices of the elements to extract.
        :return: Tuple of the new getitem nodes, corresponding to the indices.
        """
        with self.graph.inserting_after(tuple_node):
            return tuple(
                self.graph.call_function(operator.getitem, (tuple_node, idx))
                for idx in indices)

    def insert_auto_fn(self, op: OpOverload, kwargs) -> Node:
        """
        Insert an auto_functionalized node with the given op and kwargs.
        """
        return self.graph.call_function(auto_functionalized, (op, ),
                                        kwargs=kwargs)

graph property 

graph: Graph

match instance-attribute 

match = match

nodes property 

nodes: list[Node]

__init__ 

__init__(match: Match)
Source code in vllm/compilation/multi_output_match.py 
27
28
def __init__(self, match: pm.Match):
    self.match = match

find_auto_fn 

find_auto_fn(op) -> Node

Find the first auto_functionalized node with the given op in the match.

Source code in vllm/compilation/multi_output_match.py 
67
68
69
70
71
def find_auto_fn(self, op) -> fx.Node:
    """
    Find the first auto_functionalized node with the given op in the match.
    """
    return find_auto_fn(self.nodes, op)

insert_auto_fn 

insert_auto_fn(op: OpOverload, kwargs) -> Node

Insert an auto_functionalized node with the given op and kwargs.

Source code in vllm/compilation/multi_output_match.py 
104
105
106
107
108
109
def insert_auto_fn(self, op: OpOverload, kwargs) -> Node:
    """
    Insert an auto_functionalized node with the given op and kwargs.
    """
    return self.graph.call_function(auto_functionalized, (op, ),
                                    kwargs=kwargs)

insert_getitems 

insert_getitems(
    tuple_node: Node, indices: Iterable[int]
) -> tuple[Node, ...]

Insert operator.getitem nodes to extract elements from a tuple node.

:param tuple_node: The tuple node to extract elements from. :param indices: The indices of the elements to extract. :return: Tuple of the new getitem nodes, corresponding to the indices.

Source code in vllm/compilation/multi_output_match.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
def insert_getitems(self, tuple_node: fx.Node,
                    indices: Iterable[int]) -> tuple[fx.Node, ...]:
    """
    Insert operator.getitem nodes to extract elements from a tuple node.

    :param tuple_node: The tuple node to extract elements from.
    :param indices: The indices of the elements to extract.
    :return: Tuple of the new getitem nodes, corresponding to the indices.
    """
    with self.graph.inserting_after(tuple_node):
        return tuple(
            self.graph.call_function(operator.getitem, (tuple_node, idx))
            for idx in indices)

inserting_after_match 

inserting_after_match()

Insert nodes after the last node in the match. This is done to avoid use-before-definition errors after inserting replacement nodes.

Source code in vllm/compilation/multi_output_match.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
def inserting_after_match(self):
    """
    Insert nodes after the last node in the match.
    This is done to avoid use-before-definition errors after inserting
    replacement nodes.
    """

    # match.nodes is not guaranteed to be sorted.
    # Find the last node in the match.
    for last_node_in_match in reversed(self.graph.nodes):
        if last_node_in_match in self.match.nodes:
            break
    else:
        raise ValueError("No nodes in graph")

    return self.graph.inserting_after(last_node_in_match)

process abstractmethod 

process()

Process a multi-output match and manually insert the replacement.

This method should: 1. Insert the replacement nodes after the last node in the match. 2. Rebind the users of nodes in the match to use the new nodes. 3. Set meta["val"] for de-functionalization.

The result of an auto-functionalized node is a tuple of tensors. The first element is the return value of the function, usually None. The remaining elements are the mutated args of the function.

All auto-functionalized nodes must contain a proper meta["val"], as it is used by de-functionalization. meta["val"] has to contain the value of the node (tuple of tensors) that would be returned by the functionalized node during tracing.

Existing nodes in the graph all have this property set, but we have to set it manually for new nodes we insert.

Example:

op schema: foo(a: Tensor!, b: Tensor, c: Tensor!) -> None

at = auto_functionalized(torch.ops._C.foo.default, a, b, c)

at.meta["val"] = (None, a, c)

Source code in vllm/compilation/multi_output_match.py 
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
@abstractmethod
def process(self):
    """
    Process a multi-output match and manually insert the replacement.

    This method should:
    1. Insert the replacement nodes after the last node in the match.
    2. Rebind the users of nodes in the match to use the new nodes.
    3. Set meta["val"] for de-functionalization.

    The result of an auto-functionalized node is a tuple of tensors.
    The first element is the return value of the function, usually None.
    The remaining elements are the mutated args of the function.

    All auto-functionalized nodes must contain a proper meta["val"],
    as it is used by de-functionalization. meta["val"] has to contain the
    value of the node (tuple of tensors) that would be returned by the
    functionalized node during tracing.

    Existing nodes in the graph all have this property set, but we have
    to set it manually for new nodes we insert.

    Example:
    # op schema: foo(a: Tensor!, b: Tensor, c: Tensor!) -> None
    at = auto_functionalized(torch.ops._C.foo.default, a, b, c)
    # at.meta["val"] = (None, a, c)
    """
    raise NotImplementedError