[Compilation] Skip x.size(dim) in _decompose_size_nodes (#42543)
Signed-off-by: nemanjaudovic <nudovic@amd.com>
nemanjaudovic committed
Jul 16, 2026 at 02:59 UTC
81e13a05913f5365959409f6374aa624807b378a
2 files changed
+69
tests/compile/test_graph_partition.py
+64
@@ -701,3 +701,67 @@ def test_decompose_size_with_getitem_user():
701
f"getitem node '{node.name}' has {len(node.args)} args "
702
f"(expected 2): {node.args}"
703
)
704
+
705
+
706
+def test_decompose_size_leaves_scalar_size_with_dim():
707
+ """
708
+ Regression test: _decompose_size_nodes must leave x.size(dim) alone.
709
+
710
+ x.size() returns a torch.Size tuple that can't cross split boundaries and
711
+ must be decomposed. x.size(dim), however, already returns a scalar
712
+ SymInt/int that crosses fine, so the pass must not touch it.
713
+
714
+ The punica LoRA path traces token_lora_mapping[:x.size(0)] under a dynamic
715
+ batch dim, so the size(0) node ends up nested inside a slice object:
716
+
717
+ %size = call_method[target="size"](args = (%x, 0))
718
+ %slice = call_function[target=getitem](
719
+ args = (%mapping, slice(None, %size, None)))
720
+
721
+ The old pass tried to decompose this scalar node too and then erase it, but
722
+ the slice still referenced it, raising "Tried to erase Node size but it
723
+ still had N users". The fix skips size calls that carry a dim argument.
724
+ """
725
+ from torch._dynamo.source import LocalSource
726
+ from torch._subclasses.fake_tensor import FakeTensorMode
727
+ from torch.fx.experimental.symbolic_shapes import ShapeEnv
728
+
729
+ # Build graph:
730
+ # %x = placeholder
731
+ # %mapping = placeholder
732
+ # %size = x.size(0) # scalar, with a dim arg
733
+ # %sliced = mapping[slice(None, %size, None)] # size node inside a slice
734
+ graph = fx.Graph()
735
+ x = graph.placeholder("x")
736
+ mapping = graph.placeholder("token_lora_mapping")
737
+ size_node = graph.call_method("size", args=(x, 0))
738
+ sliced_node = graph.call_function(
739
+ operator.getitem,
740
+ args=(mapping, slice(None, size_node, None)),
741
+ )
742
+ graph.output((sliced_node,))
743
+
744
+ # dim 0 dynamic (SymInt) — the realistic Unsloth + LoRA case. Without the
745
+ # skip, the pass would build per-dim replacements and then crash trying to
746
+ # erase the still-referenced size node.
747
+ shape_env = ShapeEnv()
748
+ src = LocalSource("tokens")
749
+ sym_tokens = shape_env.create_symintnode(shape_env.create_symbol(4, src), hint=4)
750
+ fake_mode = FakeTensorMode(shape_env=shape_env)
751
+ with fake_mode:
752
+ fake_x = torch.empty_strided((sym_tokens, 8), (8, 1))
753
+ x.meta["example_value"] = fake_x
754
+
755
+ gm = fx.GraphModule(torch.nn.Module(), graph)
756
+
757
+ # Must not raise "Tried to erase Node ... still had N users".
758
+ _decompose_size_nodes(gm)
759
+
760
+ # The scalar x.size(0) node is left in place, untouched.
761
+ remaining = list(gm.graph.find_nodes(op="call_method", target="size"))
762
+ assert len(remaining) == 1, (
763
+ f"x.size(0) should be left untouched, found {len(remaining)} size nodes"
764
+ )
765
+ assert remaining[0].args == (x, 0), (
766
+ f"size node args changed: {remaining[0].args} (expected (x, 0))"
767
+ )
vllm/compilation/backends.py
+5
@@ -489,6 +489,11 @@ def _decompose_size_nodes(graph: fx.GraphModule) -> None:
489
size_nodes = list(graph.graph.find_nodes(op="call_method", target="size"))
490
491
for node in size_nodes:
492
+ # Only x.size() (no dim) returns a torch.Size tuple that can't cross
493
+ # split boundaries. x.size(dim) already returns a scalar SymInt/int,
494
+ # which crosses fine, so leave it untouched.
495
+ if len(node.args) > 1 or "dim" in node.kwargs:
496
+ continue
497
tensor_node = node.args[0]
498
ev = tensor_node.meta.get("example_value")
499
assert ev is not None, (