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 7, 01:30 UTC
[cudagraph] fix cudagraph warning in deepseekv32 (#28044)
Signed-off-by: zjy0516 <riverclouds.zhu@qq.com>
Jiangyun Zhu committed
Mar 9, 2026 at 08:27 UTC
e5ff140216272c529261b02b6fd13fc480713735
2 files changed
+108
tests/compile/test_graph_partition.py
+53
@@ -184,3 +184,56 @@ def test_consecutive_ops_in_split():
184
assert [node.op for node in splitting_gm.graph.nodes] == ["placeholder"] + 2 * [
185
"call_function"
186
] + ["output"]
187
+
188
+
189
+def test_empty_only_partition_is_merged():
190
+ """
191
+ Test that an empty-allocation-only partition is merged into its previous
192
+ partition during Dynamo FX splitting.
193
+ """
194
+
195
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
196
+ y = torch.sin(x)
197
+ out = torch.empty_like(y)
198
+ torch.ops.aten.cos.out(y, out=out)
199
+ return out
200
+
201
+ x = torch.randn(4, 3)
202
+ gm = make_fx(model_fn)(x)
203
+
204
+ split_ops = ["aten::sin", "aten::cos.out"]
205
+ split_gm, split_items = split_graph(gm, split_ops)
206
+
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"
210
+
211
+ output_original = gm(x)
212
+ output_split = split_gm(x)
213
+ assert torch.allclose(output_original, output_split), "Output mismatch after split"
214
+
215
+
216
+def test_builtin_empty_only_partition_is_merged():
217
+ """
218
+ In Dynamo graphs, torch.empty/empty_like may appear as builtin call targets
219
+ (not aten OpOverload). Ensure empty-only partitions are still merged.
220
+ """
221
+
222
+ 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
228
+
229
+ gm = torch.fx.symbolic_trace(model_fn)
230
+ split_gm, split_items = split_graph(gm, ["silly::attention"])
231
+
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"
235
+
236
+ x = torch.randn(2, 3, device="cuda")
237
+ output_original = gm(x)
238
+ output_split = split_gm(x)
239
+ assert torch.allclose(output_original, output_split), "Output mismatch after split"
vllm/compilation/backends.py
+55
@@ -9,6 +9,7 @@ import operator
9
import os
10
import pprint
11
import time
12
+from collections import defaultdict
13
from collections.abc import Callable, Generator, Sequence
14
from contextlib import contextmanager
15
from copy import deepcopy
@@ -405,6 +406,58 @@ class SplitItem:
406
graph: fx.GraphModule
407
408
409
+def _is_empty_allocation_node(node: fx.Node) -> bool:
410
+ if node.op == "call_method":
411
+ return node.target == "new_empty"
412
+
413
+ if node.op != "call_function":
414
+ return False
415
+
416
+ target = node.target
417
+ if target in (torch.empty, torch.empty_like, torch.empty_strided):
418
+ return True
419
+
420
+ if isinstance(target, torch._ops.OpOverloadPacket):
421
+ packet_name = target._qualified_op_name
422
+ elif isinstance(target, torch._ops.OpOverload):
423
+ packet_name = target.name()
424
+ else:
425
+ return False
426
+
427
+ return packet_name.startswith("aten::empty") or packet_name.startswith(
428
+ "aten::new_empty"
429
+ )
430
+
431
+
432
+def _merge_empty_only_subgraphs(
433
+ node_to_subgraph_id: dict[fx.Node, int],
434
+) -> None:
435
+ """
436
+ Merge a partition that only contains an empty allocation op into the
437
+ previous partition. This avoids generating standalone empty submodules,
438
+ which can lead to empty cudagraph captures.
439
+ """
440
+
441
+ 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)
446
+ nodes_by_subgraph_id[subgraph_id].append(node)
447
+
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
457
+ continue
458
+ prev_subgraph_id = subgraph_id
459
+
460
+
461
def split_graph(
462
graph: fx.GraphModule, splitting_ops: list[str]
463
) -> tuple[fx.GraphModule, list[SplitItem]]:
@@ -443,6 +496,8 @@ def split_graph(
496
else:
497
node_to_subgraph_id[node] = subgraph_id
498
499
+ _merge_empty_only_subgraphs(node_to_subgraph_id)
500
+
501
# `keep_original_order` is important!
502
# otherwise pytorch might reorder the nodes and
503
# the semantics of the graph will change when we