Skip to content

vllm.compilation.passes.ir.clone_elimination

Classes:

Functions:

UnsafeCloneEliminationPass

Bases: VllmInductorPass

This pass removes clone nodes that are no longer needed after vLLM IR lowering. It uses donated_input_ids to eliminate clones of donated graph inputs, preserving contents of non-donated graph inputs.

It is "unsafe" because it only handles aliasing of graph outputs. Solving aliasing is an open problem, so this pass intends to support known vLLM cases and not guarantee soundness on general graphs. In the future, this pass will likely support basic forms of aliasing to handle simple views (e.g. qkv -> q,k,v).

Source code in vllm/compilation/passes/ir/clone_elimination.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
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
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
class UnsafeCloneEliminationPass(VllmInductorPass):
    """This pass removes clone nodes that are no longer needed after vLLM IR lowering.
    It uses donated_input_ids to eliminate clones of donated graph inputs, preserving
    contents of non-donated graph inputs.

    It is "unsafe" because it only handles aliasing of graph outputs. Solving
    aliasing is an open problem, so this pass intends to support known vLLM cases
    and not guarantee soundness on general graphs. In the future, this pass will likely
    support basic forms of aliasing to handle simple views (e.g. qkv -> q,k,v).
    """

    def __init__(self, vllm_config: VllmConfig) -> None:
        super().__init__(vllm_config)

    @VllmInductorPass.time_and_log
    def __call__(self, graph: fx.Graph) -> None:
        count = 0
        node_to_idx = {node: i for i, node in enumerate(graph.nodes)}
        pass_context = get_pass_context()
        donated_input_ids = pass_context.donated_input_ids
        logger.debug("Donated input ids: %s", donated_input_ids)

        for node in graph.nodes:
            if not is_func(node, torch.ops.aten.clone.default):
                continue

            original_node = node.args[0]
            assert isinstance(original_node, fx.Node)

            if not clone_preserves_layout(node, original_node):
                logger.debug(
                    "Clone removal not possible, clone changes layout: "
                    "original_node=%s node=%s",
                    original_node,
                    node,
                )
                continue

            # Clone needs to be preserved if node is getting written to and
            # the old value is used again.
            # This could only happen if an inplace implementation was lowered.
            # Then node (the clone) will have one write.
            # TODO(luka) hopefully this can be removed once we lower functional graphs.
            write_idxs = [
                node_to_idx[u] for u in node.users if user_writes_to_node(u, node)
            ]
            assert len(write_idxs) in (0, 1)
            if write_idxs:
                # Check if a user of original_node occurs after a write
                write_idx = write_idxs[0]
                if any(
                    node_to_idx[orig_user] > write_idx
                    for orig_user in original_node.users
                ):
                    logger.debug(
                        "Clone removal not possible, "
                        "original_node=%s used after mutation on node=%s",
                        original_node,
                        node,
                    )
                    continue

                # Check if a node is a (non-donated) graph input
                if (
                    original_node.op == "placeholder"
                    and node_to_idx[original_node] not in donated_input_ids
                ):
                    logger.debug(
                        "Graph input %s not donated, cannot eliminate its clone",
                        original_node,
                    )
                    continue
            elif would_alias_graph_io(node, original_node):
                logger.debug(
                    "Clone removal not possible, returned clone %s would alias "
                    "graph input or output %s",
                    node,
                    original_node,
                )
                continue

            logger.debug(
                "Node %s is a redundant clone node of %s, removing it",
                node,
                original_node,
            )
            node.replace_all_uses_with(original_node)
            graph.erase_node(node)
            count += 1

        logger.debug("CloneCleanupPass removed %d clone nodes", count)

would_alias_graph_io(clone, original)

Whether replacing clone with original makes a graph output share storage with a graph input or another output, which callers may mutate in place. Unknown storage is treated as aliasing.

Source code in vllm/compilation/passes/ir/clone_elimination.py 
73
74
75
76
77
78
79
80
81
82
83
84
85
def would_alias_graph_io(clone: fx.Node, original: fx.Node) -> bool:
    """Whether replacing clone with original makes a graph output share storage
    with a graph input or another output, which callers may mutate in place.
    Unknown storage is treated as aliasing."""
    clone_storage = get_node_storage(clone)
    original_storage = get_node_storage(original)
    if clone_storage is None or original_storage is None:
        return True

    graph = clone.graph
    outputs = {get_node_storage(n) for n in graph.output_node().all_input_nodes}
    inputs = {get_node_storage(n) for n in graph.find_nodes(op="placeholder")}
    return clone_storage in outputs and original_storage in inputs | outputs