[torch.compile] fix sym_tensor_indices (vllm-project#12191)

youkaichao · Isotr0py · commit f22c8ac54da2 · 2025-02-02T21:35:02.000+08:00
Signed-off-by: youkaichao &lt;youkaichao@gmail.com&gt;
Signed-off-by: Isotr0py &lt;2037008807@qq.com&gt;
diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py
@@ -624,9 +624,13 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:
         ]
 
         # index of tensors that have symbolic shapes (batch size)
+        # for weights and static buffers, they will have concrete shapes.
+        # symbolic shape only happens for input tensors.
+        from torch.fx.experimental.symbolic_shapes import is_symbolic
         self.sym_tensor_indices = [
             i for i, x in enumerate(fake_args)
-            if isinstance(x, torch._subclasses.fake_tensor.FakeTensor)
+            if isinstance(x, torch._subclasses.fake_tensor.FakeTensor) and \
+                any(is_symbolic(d) for d in x.size())
         ]
 
         # compiler managed cudagraph input buffers

Original file line number	Diff line number	Diff line change
`@@ -624,9 +624,13 @@ def __call__(self, graph: fx.GraphModule, example_inputs) -> Callable:`
`624`	`624`	`]`
`625`	`625`
`626`	`626`	`# index of tensors that have symbolic shapes (batch size)`
	`627`	`+ # for weights and static buffers, they will have concrete shapes.`
	`628`	`+ # symbolic shape only happens for input tensors.`
	`629`	`+ from torch.fx.experimental.symbolic_shapes import is_symbolic`
`627`	`630`	`self.sym_tensor_indices = [`
`628`	`631`	`i for i, x in enumerate(fake_args)`
`629`		`- if isinstance(x, torch._subclasses.fake_tensor.FakeTensor)`
	`632`	`+ if isinstance(x, torch._subclasses.fake_tensor.FakeTensor) and \`
	`633`	`+ any(is_symbolic(d) for d in x.size())`
`630`	`634`	`]`
`631`	`635`
`632`	`636`	`# compiler managed cudagraph input buffers`