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 5, 06:00 UTC
[compile] Bug fix for _decompose_size_nodes (#38360)
Signed-off-by: Animesh Jain <anijain@umich.edu>
Animesh Jain committed
Apr 12, 2026 at 13:20 UTC
f00c5539d7265719c188a5335159f68178f66e2f
2 files changed
+99
-10
tests/compile/test_graph_partition.py
+75
-1
@@ -9,7 +9,11 @@ import torch._dynamo
9
import torch.fx as fx
10
from torch.fx.experimental.proxy_tensor import make_fx
11
12
-from vllm.compilation.backends import _is_empty_allocation_node, split_graph
12
+from vllm.compilation.backends import (
13
+ _decompose_size_nodes,
14
+ _is_empty_allocation_node,
15
+ split_graph,
16
+)
17
from vllm.compilation.passes.fx_utils import find_op_nodes
18
19
# This import automatically registers `torch.ops.silly.attention`
@@ -622,3 +626,73 @@ def test_sym_size_metadata_propagated():
626
else:
627
example_inputs.append(int(ev))
628
standalone_compile(submod, example_inputs, dynamic_shapes="from_example_inputs")
629
+
630
+
631
+def test_decompose_size_with_getitem_user():
632
+ """
633
+ Regression test: _decompose_size_nodes must handle getitem users of size()
634
+ correctly.
635
+
636
+ When a graph contains x.shape[i], it can appear as:
637
+
638
+ %size = call_method[target="size"](args = (%x,))
639
+ %getitem = call_function[target=operator.getitem](args = (%size, 1))
640
+
641
+ The old code spliced *all* per-dim values into every user's args
642
+ unconditionally, turning the 2-arg getitem into a malformed 3-arg node:
643
+
644
+ %getitem(args = (%sym_size_int, 5120, 1)) # TypeError at runtime
645
+
646
+ The fix detects getitem users and replaces them with dims[idx] directly.
647
+ """
648
+ # Build a graph manually to guarantee the size() + getitem pattern.
649
+ #
650
+ # Graph:
651
+ # %x = placeholder
652
+ # %size = x.size()
653
+ # %dim1 = getitem(%size, 1) <-- the getitem branch we're testing
654
+ # %relu = relu(%x)
655
+ # %view = view(%relu, -1, %dim1)
656
+ # return %view
657
+ graph = fx.Graph()
658
+ x = graph.placeholder("x")
659
+ size_node = graph.call_method("size", args=(x,))
660
+ getitem_node = graph.call_function(operator.getitem, args=(size_node, 1))
661
+ relu_node = graph.call_function(torch.ops.aten.relu.default, args=(x,))
662
+ view_node = graph.call_function(
663
+ torch.ops.aten.view.default, args=(relu_node, [-1, getitem_node])
664
+ )
665
+ graph.output(view_node)
666
+
667
+ # Attach example_value metadata so _decompose_size_nodes can inspect dims.
668
+ # dim 0 is dynamic (SymInt), dim 1 is static (8).
669
+ from torch._dynamo.source import LocalSource
670
+ from torch._subclasses.fake_tensor import FakeTensorMode
671
+ from torch.fx.experimental.symbolic_shapes import ShapeEnv
672
+
673
+ shape_env = ShapeEnv()
674
+ src = LocalSource("batch_size")
675
+ sym_batch = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4)
676
+ fake_mode = FakeTensorMode(shape_env=shape_env)
677
+ with fake_mode:
678
+ fake_x = torch.empty_strided((sym_batch, 8), (8, 1))
679
+ x.meta["example_value"] = fake_x
680
+
681
+ gm = fx.GraphModule(torch.nn.Module(), graph)
682
+
683
+ # Run decomposition — this would produce a 3-arg getitem without the fix
684
+ _decompose_size_nodes(gm)
685
+
686
+ # Verify no size() nodes remain
687
+ remaining_size_nodes = list(gm.graph.find_nodes(op="call_method", target="size"))
688
+ assert len(remaining_size_nodes) == 0, (
689
+ f"size() nodes should be fully decomposed, found {len(remaining_size_nodes)}"
690
+ )
691
+
692
+ # Verify no malformed getitem nodes (3+ args)
693
+ for node in gm.graph.nodes:
694
+ if node.op == "call_function" and node.target is operator.getitem:
695
+ assert len(node.args) == 2, (
696
+ f"getitem node '{node.name}' has {len(node.args)} args "
697
+ f"(expected 2): {node.args}"
698
+ )
vllm/compilation/backends.py
+24
-9
@@ -516,16 +516,31 @@ def _decompose_size_nodes(graph: fx.GraphModule) -> None:
516
)
517
518
# Replace size node in each user's args.
519
- # Dynamo always passes size as a direct arg: view(clone, size)
520
- # → view(clone, d0, d1, ...)
519
for user in list(node.users):
522
- new_args = []
523
- for arg in user.args:
524
- if arg is node:
525
- new_args.extend(dims)
526
- else:
527
- new_args.append(arg)
528
- user.args = tuple(new_args)
520
+ if (
521
+ user.op == "call_function"
522
+ and user.target is operator.getitem
523
+ and len(user.args) == 2
524
+ and user.args[0] is node
525
+ ):
526
+ # getitem(size, idx) → replace with dims[idx] directly.
527
+ idx = user.args[1]
528
+ assert isinstance(idx, int), (
529
+ f"Expected literal int index for getitem on size(), "
530
+ f"got {type(idx).__name__}: {idx}"
531
+ )
532
+ user.replace_all_uses_with(dims[idx])
533
+ graph.graph.erase_node(user)
534
+ else:
535
+ # User consumes the full size tuple (e.g. view(clone, size))
536
+ # → view(clone, d0, d1, ...)
537
+ new_args = []
538
+ for arg in user.args:
539
+ if arg is node:
540
+ new_args.extend(dims)
541
+ else:
542
+ new_args.append(arg)
543
+ user.args = tuple(new_args)
544
graph.graph.erase_node(node)
545
546