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 8, 12:30 UTC
[torch.compile] Stop doing unnecessary FakeTensorProp in PiecewiseCompileInterpreter (#34093)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Feb 10, 2026 at 22:15 UTC
e30cedd44be332e1ddc7ec43b8a33bce532e7614
2 files changed
+41
-4
tests/compile/fullgraph/test_simple.py
+38
-3
@@ -27,10 +27,29 @@ from ...utils import create_new_process_for_each_test
27
from ..silly_attention import get_global_counter, reset_global_counter
28
29
30
+# Custom op that returns an unbacked symint during graph capture
31
+@torch.library.custom_op("mylib::foo", mutates_args=())
32
+def foo(x: torch.Tensor) -> int:
33
+ return 3
34
+
35
+
36
+@foo.register_fake
37
+def _(x):
38
+ return torch.library.get_ctx().new_dynamic_size()
39
+
40
+
41
@support_torch_compile
42
class SillyModel(nn.Module):
32
- def __init__(self, *, vllm_config: VllmConfig, prefix: str = "", **kwargs) -> None:
43
+ def __init__(
44
+ self,
45
+ *,
46
+ vllm_config: VllmConfig,
47
+ prefix: str = "",
48
+ intermediate_unbacked=False,
49
+ **kwargs,
50
+ ) -> None:
51
super().__init__()
52
+ self.intermediate_unbacked = intermediate_unbacked
53
54
def forward(self, x: torch.Tensor) -> torch.Tensor:
55
"""
@@ -44,6 +63,13 @@ class SillyModel(nn.Module):
63
torch.ops.silly.attention(x, x, x, out)
64
x = out
65
x = x - 2
66
+
67
+ if self.intermediate_unbacked:
68
+ # Test for unbacked symints: the following is a fancy way to multiply by 1
69
+ u0 = foo(x)
70
+ ones = x.new_ones(x.shape[0], u0).sum(-1) / 3
71
+ x = x * ones
72
+
73
x = x - 1
74
out = torch.empty_like(x)
75
torch.ops.silly.attention(x, x, x, out)
@@ -52,6 +78,7 @@ class SillyModel(nn.Module):
78
return x
79
80
81
+@torch._dynamo.config.patch(capture_dynamic_output_shape_ops=True)
82
def _run_simple_model(
83
splitting_ops,
84
use_inductor_graph_partition,
@@ -60,6 +87,8 @@ def _run_simple_model(
87
expected_num_piecewise_capturable_graphs_seen,
88
expected_num_backend_compilations,
89
expected_num_cudagraph_captured,
90
+ *,
91
+ intermediate_unbacked=False,
92
):
93
vllm_config = VllmConfig(
94
compilation_config=CompilationConfig(
@@ -72,7 +101,11 @@ def _run_simple_model(
101
)
102
)
103
with set_current_vllm_config(vllm_config):
75
- model = SillyModel(vllm_config=vllm_config, prefix="")
104
+ model = SillyModel(
105
+ vllm_config=vllm_config,
106
+ prefix="",
107
+ intermediate_unbacked=intermediate_unbacked,
108
+ )
109
110
inputs = torch.randn(100).cuda()
111
@@ -125,9 +158,10 @@ def _run_simple_model(
158
159
160
@pytest.mark.parametrize("backend", ["inductor", "eager"])
161
+@pytest.mark.parametrize("intermediate_unbacked", [True, False])
162
@torch.inference_mode()
163
@create_new_process_for_each_test("spawn")
130
-def test_simple_piecewise_compile(backend):
164
+def test_simple_piecewise_compile(backend, intermediate_unbacked):
165
_run_simple_model(
166
splitting_ops=["silly::attention"],
167
use_inductor_graph_partition=False,
@@ -140,6 +174,7 @@ def test_simple_piecewise_compile(backend):
174
expected_num_backend_compilations=3,
175
# num_cudagraph_sizes * num_piecewise_capturable_graphs_seen
176
expected_num_cudagraph_captured=6,
177
+ intermediate_unbacked=intermediate_unbacked,
178
)
179
180
vllm/compilation/backends.py
+3
-1
@@ -570,7 +570,9 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter): # type: ignore[misc]
570
) -> Any:
571
assert isinstance(target, str)
572
573
- output = super().call_module(target, args, kwargs)
573
+ gm = getattr(self.module, target)
574
+ outputs = gm.graph.output_node().args[0]
575
+ output = fx.map_arg(outputs, lambda node: node.meta["example_value"])
576
577
if target in self.compile_submod_names:
578
index = self.compile_submod_names.index(target)