A high-throughput and memory-efficient inference and serving engine for LLMs
Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git
Synced Nov 1, 20:00 UTC
Preserve layout-changing clones (#44574)
Signed-off-by: Michael Gschwind <mgschwind@nvidia.com> Co-authored-by: Michael Gschwind <mgschwind@nvidia.com>
Mike G committed
Jun 5, 2026 at 17:45 UTC
2f27c9a1500accc4faf7217f2af254a981df084e
2 files changed
+48
tests/compile/passes/ir/test_clone_cleanup.py
+19
@@ -132,6 +132,25 @@ class TestCloneCleanup:
132
assert count_clones(graph_module.graph) == 0
133
torch.testing.assert_close(actual, expected)
134
135
+ def test_keep_clone_that_changes_layout(self, clone_cleanup_pass):
136
+ """Clone must be kept when it materializes a compact slice layout."""
137
+
138
+ def f(x: torch.Tensor) -> torch.Tensor:
139
+ return x[:, :3].contiguous()
140
+
141
+ inp = torch.randn(4, 5)
142
+ graph_module = make_fx(f)(inp)
143
+ assert count_clones(graph_module.graph) == 1
144
+
145
+ expected = graph_module(inp)
146
+ clone_cleanup_pass(graph_module.graph)
147
+ graph_module.recompile()
148
+ actual = graph_module(inp)
149
+
150
+ assert count_clones(graph_module.graph) == 1
151
+ assert actual.stride() == expected.stride() == (3, 1)
152
+ torch.testing.assert_close(actual, expected)
153
+
154
def test_multiple_clones_of_same_input(self, clone_cleanup_pass):
155
"""Test multiple independent clones of the same input."""
156
vllm/compilation/passes/ir/clone_elimination.py
+29
@@ -16,6 +16,26 @@ from ..vllm_inductor_pass import VllmInductorPass
16
logger = init_logger(__name__)
17
18
19
+def clone_preserves_layout(node: fx.Node, original_node: fx.Node) -> bool:
20
+ node_val = node.meta.get("val")
21
+ original_val = original_node.meta.get("val")
22
+ if node_val is None or original_val is None:
23
+ return True
24
+
25
+ try:
26
+ node_stride = tuple(node_val.stride())
27
+ original_stride = tuple(original_val.stride())
28
+ node_storage_offset = node_val.storage_offset()
29
+ original_storage_offset = original_val.storage_offset()
30
+ except (AttributeError, RuntimeError):
31
+ return True
32
+
33
+ return (
34
+ node_stride == original_stride
35
+ and node_storage_offset == original_storage_offset
36
+ )
37
+
38
+
39
def user_writes_to_node(user: fx.Node, node: fx.Node) -> bool:
40
if user.op == "output":
41
return False
@@ -79,6 +99,15 @@ class UnsafeCloneEliminationPass(VllmInductorPass):
99
original_node = node.args[0]
100
assert isinstance(original_node, fx.Node)
101
102
+ if not clone_preserves_layout(node, original_node):
103
+ logger.debug(
104
+ "Clone removal not possible, clone changes layout: "
105
+ "original_node=%s node=%s",
106
+ original_node,
107
+ node,
108
+ )
109
+ continue
110
+
111
# Clone needs to be preserved if node is getting written to and
112
# the old value is used again.
113
# This could only happen if an inplace implementation was lowered.