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 14, 05:30 UTC
[BugFix] Fix noop elimination edge case (#26394)
Signed-off-by: Andy Lo <andy@mistral.ai>
Andy Lo committed
Oct 10, 2025 at 14:33 UTC
67661375fad8dbea1d9f3b812c55ec42758cf7d2
2 files changed
+40
-51
tests/compile/test_noop_elimination.py
+12
-3
@@ -12,15 +12,23 @@ from .backend import TestBackend
12
13
14
@pytest.mark.parametrize("dtype", [torch.float16, torch.bfloat16, torch.float32])
15
-@pytest.mark.parametrize("num_tokens", [256, 1024])
15
+# Important edge case is when `num_tokens == buffer_size`
16
+@pytest.mark.parametrize(
17
+ ("num_tokens", "buffer_size"), [(256, 256), (256, 512), (1024, 1024), (1024, 1025)]
18
+)
19
@pytest.mark.parametrize("hidden_size", [64, 4096])
17
-def test_noop_elimination(dtype, num_tokens, hidden_size):
20
+def test_noop_elimination(dtype, num_tokens, hidden_size, buffer_size):
21
torch.set_default_device("cuda")
22
torch.set_default_dtype(dtype)
23
torch.manual_seed(1)
24
25
class Model(torch.nn.Module):
26
+ def __init__(self) -> None:
27
+ super().__init__()
28
+ self.pos_embed = torch.empty(buffer_size, hidden_size, dtype=dtype)
29
+
30
def forward(self, x):
31
+ x += self.pos_embed[: x.shape[0]]
32
# Chain of reshapes
33
y = x.reshape(-1, 128, 32)
34
z = y.reshape(-1, 4096)
@@ -65,9 +73,10 @@ def test_noop_elimination(dtype, num_tokens, hidden_size):
73
torch.testing.assert_close(result, result2, atol=ATOL, rtol=RTOL)
74
75
# The no-op reshape and slice should be eliminated.
76
+ # The initial slice on the positional embedding should remain.
77
# The chain of reshapes should be fused into a single reshape.
78
assert backend.op_count(torch.ops.aten.reshape.default) == 1
70
- assert backend.op_count(torch.ops.aten.slice.Tensor) == 0
79
+ assert backend.op_count(torch.ops.aten.slice.Tensor) == 1
80
assert backend.op_count(torch.ops.aten.slice_scatter.default) == 0
81
82
vllm/compilation/noop_elimination.py
+28
-48
@@ -81,49 +81,32 @@ class NoOpEliminationPass(VllmInductorPass):
81
graph.erase_node(input)
82
count += 1
83
84
- # Case 2: remove this reshape if it produces the original shape
85
- input, shape = node.args[:2]
86
- input_shape = input.meta["val"].shape
87
- if len(shape) != len(input_shape):
88
- # Reshape changing rank, skip
89
- continue
90
-
91
- if shape.count(-1) > 1:
92
- # Invalid reshape args, skip
93
- continue
94
-
95
- if self.reshape_all_dims_equivalent(shape, input_shape):
96
- node.replace_all_uses_with(input)
97
- graph.erase_node(node)
98
- count += 1
99
-
100
- elif is_func(node, torch.ops.aten.slice.Tensor):
101
- # python slicing semantics are different from reshape
102
- # Don't treat -1 as inferred dimension
103
- input, dim_index, start, end = node.args[:4]
84
+ # remove reshape/slice if it produces the original shape
85
+ if is_func(node, torch.ops.aten.reshape.default) or is_func(
86
+ node, torch.ops.aten.slice.Tensor
87
+ ):
88
+ input = node.args[0]
89
input_shape = input.meta["val"].shape
90
output_shape = node.meta["val"].shape
106
-
107
- if output_shape == input_shape:
91
+ if self.all_dims_equivalent(input_shape, output_shape):
92
node.replace_all_uses_with(input)
93
graph.erase_node(node)
94
count += 1
111
-
95
elif is_func(node, torch.ops.aten.slice_scatter.default):
96
base, view, dim_index, start, end = node.args[:5]
97
base_shape = base.meta["val"].shape
98
view_shape = view.meta["val"].shape
99
117
- if base_shape == view_shape:
100
+ if self.all_dims_equivalent(base_shape, view_shape):
101
node.replace_all_uses_with(view)
102
graph.erase_node(node)
103
count += 1
104
105
logger.debug("Removed %s no-op reshapes and slices", count)
106
124
- # ---------------------- Reshape helpers ----------------------
125
- def reshape_dims_equivalent(
126
- self, dim: Union[int, torch.fx.Node], i_dim: Union[int, SymInt]
107
+ # ---------------------- Shape comparison helpers ----------------------
108
+ def dims_equivalent(
109
+ self, dim: Union[int, SymInt], i_dim: Union[int, SymInt]
110
) -> bool:
111
"""
112
This function checks if two dimensions are equivalent.
@@ -131,27 +114,24 @@ class NoOpEliminationPass(VllmInductorPass):
114
:param i_dim: The corresponding dimension in the input tensor
115
:return: Are the dimensions equivalent?
116
134
- There are three cases in which the dimensions are equivalent:
117
+ There are two cases in which the dimensions are equivalent:
118
1. The dimensions are equal (both integers)
136
- 2. The reshape dimension is -1 (i.e. inferred)
137
- 3. The dimensions both correspond to the same SymInt
138
-
139
- While case 2 does not guarantee the dimensions are equal,
140
- they are equal if all other dimensions are equal.
141
-
142
- In case 3, the reshape dimension is a torch.fx.Node,
143
- and its value is a SymInt. That value is equal to the
144
- input dimension.
119
+ 2. The dimensions both correspond to the same SymInt
120
"""
146
- # Case 1 and 2
147
- if dim == i_dim or dim == -1:
148
- return True
149
- # Case 3
150
- return isinstance(dim, torch.fx.Node) and dim.meta["val"] == i_dim
151
-
152
- def reshape_all_dims_equivalent(
153
- self,
154
- dims: Iterable[Union[int, torch.fx.Node]],
155
- i_dims: Iterable[Union[int, SymInt]],
121
+ # Case 1
122
+ if isinstance(i_dim, int) and isinstance(dim, int):
123
+ return dim == i_dim
124
+ # Case 2
125
+ if isinstance(i_dim, SymInt) and isinstance(dim, SymInt):
126
+ return dim == i_dim
127
+ return False
128
+
129
+ def all_dims_equivalent(
130
+ self, dims: Iterable[Union[int, SymInt]], i_dims: Iterable[Union[int, SymInt]]
131
) -> bool:
157
- return all(self.reshape_dims_equivalent(s, i_s) for s, i_s in zip(dims, i_dims))
132
+ dims_ = list(dims)
133
+ i_dims_ = list(i_dims)
134
+ if len(dims_) != len(i_dims_):
135
+ # Different ranks can't be equivalent
136
+ return False
137
+ return all(self.dims_equivalent(s, i_s) for s, i_s in zip(dims, i_dims))