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 6, 23:00 UTC
[Bugfix] Avoid merging empty-only partitions into splitting-op subgraphs (#36595)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Mar 10, 2026 at 22:39 UTC
ca5fb4bbd85244fafba72fb91523c657025998a3
2 files changed
+132
-29
tests/compile/test_graph_partition.py
+105
-15
@@ -7,7 +7,7 @@ import pytest
7
import torch
8
from torch.fx.experimental.proxy_tensor import make_fx
9
10
-from vllm.compilation.backends import split_graph
10
+from vllm.compilation.backends import _is_empty_allocation_node, split_graph
11
from vllm.compilation.passes.fx_utils import find_op_nodes
12
13
# This import automatically registers `torch.ops.silly.attention`
@@ -186,10 +186,25 @@ def test_consecutive_ops_in_split():
186
] + ["output"]
187
188
189
-def test_empty_only_partition_is_merged():
189
+def _get_empty_nodes(split_item):
190
+ return [
191
+ node for node in split_item.graph.graph.nodes if _is_empty_allocation_node(node)
192
+ ]
193
+
194
+
195
+def _subgraphs_with_empty_nodes(split_items, *, is_splitting_graph):
196
+ return [
197
+ split_item
198
+ for split_item in split_items
199
+ if split_item.is_splitting_graph == is_splitting_graph
200
+ and _get_empty_nodes(split_item)
201
+ ]
202
+
203
+
204
+def test_empty_only_partition_stays_separate_after_splitting_predecessor():
205
"""
191
- Test that an empty-allocation-only partition is merged into its previous
192
- partition during Dynamo FX splitting.
206
+ Empty-only subgraphs should not be merged when the only predecessor is
207
+ a splitting-op subgraph.
208
"""
209
210
def model_fn(x: torch.Tensor) -> torch.Tensor:
@@ -204,9 +219,65 @@ def test_empty_only_partition_is_merged():
219
split_ops = ["aten::sin", "aten::cos.out"]
220
split_gm, split_items = split_graph(gm, split_ops)
221
207
- # Without the merge, this graph is split into 3 partitions where the
208
- # middle partition contains only aten::empty_like.
209
- assert len(split_items) == 2, "Empty-only partition should be merged"
222
+ # Graph partitioning for this pattern is:
223
+ # [sin], [empty_like], [cos.out].
224
+ assert len(split_items) == 3, (
225
+ "Empty-only partition should not merge into splitting-op subgraph"
226
+ )
227
+
228
+ splitting_with_empty = _subgraphs_with_empty_nodes(
229
+ split_items, is_splitting_graph=True
230
+ )
231
+ assert len(splitting_with_empty) == 0, (
232
+ "Splitting-op subgraphs should not contain empty allocation nodes: "
233
+ f"{[item.submod_name for item in splitting_with_empty]}"
234
+ )
235
+
236
+ output_original = gm(x)
237
+ output_split = split_gm(x)
238
+ assert torch.allclose(output_original, output_split), "Output mismatch after split"
239
+
240
+
241
+def test_empty_only_partition_is_merged():
242
+ """
243
+ Empty-only subgraphs should still be merged when a non-splitting predecessor
244
+ exists. The merged empty node must remain outside splitting-op subgraphs.
245
+ """
246
+
247
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
248
+ base = x + 1
249
+ y = torch.sin(base)
250
+ out = torch.empty_like(base)
251
+ torch.ops.aten.cos.out(base, out=out)
252
+ return out + y
253
+
254
+ x = torch.randn(4, 3)
255
+ gm = make_fx(model_fn)(x)
256
+ split_gm, split_items = split_graph(gm, ["aten::sin", "aten::cos.out"])
257
+
258
+ # Partitioning should be:
259
+ # [add, empty_like], [sin], [cos.out], [add].
260
+ assert len(split_items) == 4, (
261
+ "Empty-only partition should be merged into non-splitting predecessor"
262
+ )
263
+
264
+ splitting_with_empty = _subgraphs_with_empty_nodes(
265
+ split_items, is_splitting_graph=True
266
+ )
267
+ assert len(splitting_with_empty) == 0, (
268
+ "Splitting-op subgraphs should not contain empty allocation nodes: "
269
+ f"{[item.submod_name for item in splitting_with_empty]}"
270
+ )
271
+
272
+ non_splitting_with_empty = _subgraphs_with_empty_nodes(
273
+ split_items, is_splitting_graph=False
274
+ )
275
+ assert len(non_splitting_with_empty) == 1, (
276
+ "Exactly one non-splitting subgraph should contain the merged empty node"
277
+ )
278
+ assert len(_get_empty_nodes(non_splitting_with_empty[0])) == 1, (
279
+ "Expected exactly one empty allocation node in merged subgraph"
280
+ )
281
282
output_original = gm(x)
283
output_split = split_gm(x)
@@ -220,18 +291,37 @@ def test_builtin_empty_only_partition_is_merged():
291
"""
292
293
def model_fn(x: torch.Tensor) -> torch.Tensor:
223
- out1 = torch.empty_like(x)
224
- torch.ops.silly.attention(x, x, x, out1)
225
- out2 = torch.empty_like(x)
226
- torch.ops.silly.attention(out1, out1, out1, out2)
227
- return out2
294
+ hidden = x + 1
295
+ out1 = torch.empty_like(hidden)
296
+ torch.ops.silly.attention(hidden, hidden, hidden, out1)
297
+ out2 = torch.empty_like(hidden)
298
+ torch.ops.silly.attention(out1, out1, hidden, out2)
299
+ return out2 + hidden
300
301
gm = torch.fx.symbolic_trace(model_fn)
302
split_gm, split_items = split_graph(gm, ["silly::attention"])
303
232
- # Without the empty-only merge, this graph creates 4 partitions:
233
- # [empty_like], [attention], [empty_like], [attention].
234
- assert len(split_items) == 3, "Builtin empty-only partition should be merged"
304
+ # Without empty-only merge, this graph would split into:
305
+ # [add, empty_like], [attention], [empty_like], [attention], [add].
306
+ assert len(split_items) == 4, "Builtin empty-only partition should be merged"
307
+
308
+ splitting_with_empty = _subgraphs_with_empty_nodes(
309
+ split_items, is_splitting_graph=True
310
+ )
311
+ assert len(splitting_with_empty) == 0, (
312
+ "Splitting-op subgraphs should not contain empty allocation nodes: "
313
+ f"{[item.submod_name for item in splitting_with_empty]}"
314
+ )
315
+
316
+ non_splitting_with_empty = _subgraphs_with_empty_nodes(
317
+ split_items, is_splitting_graph=False
318
+ )
319
+ assert len(non_splitting_with_empty) == 1, (
320
+ "Exactly one non-splitting subgraph should contain merged empty nodes"
321
+ )
322
+ assert len(_get_empty_nodes(non_splitting_with_empty[0])) == 2, (
323
+ "Expected two builtin empty_like nodes in merged non-splitting subgraph"
324
+ )
325
326
x = torch.randn(2, 3, device="cuda")
327
output_original = gm(x)
vllm/compilation/backends.py
+27
-14
@@ -431,6 +431,7 @@ def _is_empty_allocation_node(node: fx.Node) -> bool:
431
432
def _merge_empty_only_subgraphs(
433
node_to_subgraph_id: dict[fx.Node, int],
434
+ split_op_graphs: list[int],
435
) -> None:
436
"""
437
Merge a partition that only contains an empty allocation op into the
@@ -439,23 +440,35 @@ def _merge_empty_only_subgraphs(
440
"""
441
442
nodes_by_subgraph_id: dict[int, list[fx.Node]] = defaultdict(list)
442
- subgraph_id_order: list[int] = []
443
for node, subgraph_id in node_to_subgraph_id.items():
444
- if subgraph_id not in nodes_by_subgraph_id:
445
- subgraph_id_order.append(subgraph_id)
444
nodes_by_subgraph_id[subgraph_id].append(node)
445
448
- prev_subgraph_id: int | None = None
449
- for subgraph_id in subgraph_id_order:
450
- nodes = nodes_by_subgraph_id[subgraph_id]
451
- if (
452
- len(nodes) == 1
453
- and _is_empty_allocation_node(nodes[0])
454
- and prev_subgraph_id is not None
455
- ):
456
- node_to_subgraph_id[nodes[0]] = prev_subgraph_id
446
+ splitting_subgraphs = set(split_op_graphs)
447
+ prev_non_splitting_subgraph_id: int | None = None
448
+
449
+ max_subgraph_id = max(node_to_subgraph_id.values(), default=-1)
450
+ for subgraph_id in range(max_subgraph_id + 1):
451
+ nodes = nodes_by_subgraph_id.get(subgraph_id, [])
452
+ if not nodes:
453
continue
458
- prev_subgraph_id = subgraph_id
454
+
455
+ is_non_splitting_subgraph = subgraph_id not in splitting_subgraphs
456
+ is_empty_only_subgraph = len(nodes) == 1 and _is_empty_allocation_node(nodes[0])
457
+ merged = False
458
+
459
+ if is_empty_only_subgraph and prev_non_splitting_subgraph_id is not None:
460
+ # Safety check: don't move allocation before any input producer.
461
+ empty_node = nodes[0]
462
+ if all(
463
+ input_node.op == "placeholder"
464
+ or node_to_subgraph_id[input_node] <= prev_non_splitting_subgraph_id
465
+ for input_node in empty_node.all_input_nodes
466
+ ):
467
+ node_to_subgraph_id[empty_node] = prev_non_splitting_subgraph_id
468
+ merged = True
469
+
470
+ if not merged and is_non_splitting_subgraph:
471
+ prev_non_splitting_subgraph_id = subgraph_id
472
473
474
def split_graph(
@@ -496,7 +509,7 @@ def split_graph(
509
else:
510
node_to_subgraph_id[node] = subgraph_id
511
499
- _merge_empty_only_subgraphs(node_to_subgraph_id)
512
+ _merge_empty_only_subgraphs(node_to_subgraph_id, split_op_graphs)
513
514
# `keep_original_order` is important!
515
# otherwise pytorch might reorder the nodes and