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 6, 21:45 UTC
[torch.compile] Use FakeTensors instead of real GPU tensors for single-size compilation (#36093)
Signed-off-by: Richard Zou <zou3519@gmail.com>
Richard Zou committed
Mar 11, 2026 at 12:07 UTC
822e250ab74899af4bc28aa5d738ec4c0e8c646e
3 files changed
+137
-25
tests/compile/test_compile_ranges.py
+82
@@ -127,6 +127,88 @@ def test_compile_config_get_compile_ranges():
127
]
128
129
130
+class PostGradStaticShapeChecker(InductorPass):
131
+ """Asserts that compile_sizes entries produce graphs with fully concrete
132
+ (non-symbolic) shapes, and compile_ranges entries have symbolic shapes."""
133
+
134
+ def __init__(self):
135
+ self.num_static_calls = 0
136
+ self.num_dynamic_calls = 0
137
+
138
+ def __call__(self, graph: fx.Graph):
139
+ from torch.fx.experimental.symbolic_shapes import is_symbolic
140
+
141
+ compile_range = get_pass_context().compile_range
142
+ is_single = compile_range.is_single_size()
143
+
144
+ for node in graph.nodes:
145
+ val = node.meta.get("val")
146
+ if val is None:
147
+ val = node.meta.get("example_value")
148
+ if isinstance(val, torch.Tensor):
149
+ has_symbolic = any(is_symbolic(d) for d in val.shape)
150
+ if is_single:
151
+ assert not has_symbolic, (
152
+ f"compile_sizes entry {compile_range}: "
153
+ f"node '{node.name}' has symbolic shape "
154
+ f"{val.shape}"
155
+ )
156
+ else:
157
+ # compile_ranges should have at least some
158
+ # symbolic shapes (the batch dimension)
159
+ if has_symbolic:
160
+ self.num_dynamic_calls += 1
161
+ return
162
+
163
+ if is_single:
164
+ self.num_static_calls += 1
165
+
166
+ def uuid(self) -> str:
167
+ state: dict[str, Any] = {}
168
+ return InductorPass.hash_dict(state)
169
+
170
+
171
+def test_compile_sizes_produce_static_shapes(use_fresh_inductor_cache):
172
+ """Verify that compile_sizes entries are compiled with fully concrete
173
+ shapes (no SymInts), while compile_ranges entries retain dynamic shapes."""
174
+ checker = PostGradStaticShapeChecker()
175
+ torch.set_default_device("cuda")
176
+ vllm_config = VllmConfig(
177
+ scheduler_config=SchedulerConfig(
178
+ max_num_batched_tokens=8192,
179
+ max_model_len=8192,
180
+ is_encoder_decoder=False,
181
+ ),
182
+ compilation_config=CompilationConfig(
183
+ mode=CompilationMode.VLLM_COMPILE,
184
+ compile_ranges_endpoints=[8],
185
+ compile_sizes=[16],
186
+ inductor_compile_config={
187
+ "post_grad_custom_post_pass": checker,
188
+ },
189
+ ),
190
+ )
191
+
192
+ with set_current_vllm_config(vllm_config):
193
+ model = TestModel(vllm_config=vllm_config, prefix="").eval()
194
+ # 3 compilations: Range(1,8), Range(9,8192), single-size 16
195
+ with compilation_counter.expect(
196
+ num_graphs_seen=1,
197
+ num_piecewise_graphs_seen=1,
198
+ num_backend_compilations=3,
199
+ ):
200
+ run_model(vllm_config, model, [1, 16, 64])
201
+
202
+ # compile_sizes=16 should produce static shapes
203
+ assert checker.num_static_calls == 1, (
204
+ f"Expected 1 static compilation, got {checker.num_static_calls}"
205
+ )
206
+ # compile_ranges should produce dynamic shapes
207
+ assert checker.num_dynamic_calls == 2, (
208
+ f"Expected 2 dynamic compilations, got {checker.num_dynamic_calls}"
209
+ )
210
+
211
+
212
def test_inductor_cache_compile_ranges(monkeypatch, use_fresh_inductor_cache):
213
# To force multiple compilations, we disable the compile cache
214
monkeypatch.setenv("VLLM_DISABLE_COMPILE_CACHE", "1")
vllm/compilation/compiler_interface.py
+29
-3
@@ -348,13 +348,39 @@ class InductorStandaloneAdaptor(CompilerInterface):
348
# Can remove this after the following issue gets fixed
349
# https://github.com/pytorch/pytorch/issues/174502
350
if envs.VLLM_ENABLE_PREGRAD_PASSES:
351
- ctx: Any = contextlib.nullcontext()
351
+ pregrad_ctx: Any = contextlib.nullcontext()
352
else:
353
- ctx = patch(
353
+ pregrad_ctx = patch(
354
"torch._inductor.compile_fx._recursive_pre_grad_passes",
355
lambda gm, _: gm,
356
)
357
- with ctx, _patch_constrain_to_fx_strides():
357
+
358
+ # When inputs are FakeTensors (from create_concrete_args),
359
+ # standalone_compile("from_example_inputs") would normally create
360
+ # a fresh FakeTensorMode, causing a mode mismatch assertion.
361
+ # Patch FakeTensorMode in standalone_compile so it reuses the
362
+ # mode already attached to our FakeTensors. This gives us both
363
+ # ignore_shape_env=True (from "from_example_inputs") and mode
364
+ # consistency (from reusing our mode).
365
+ # Can remove this after the following issue gets fixed:
366
+ # https://github.com/pytorch/pytorch/issues/176562
367
+ from torch._subclasses.fake_tensor import FakeTensor
368
+
369
+ input_fake_mode = None
370
+ for x in example_inputs:
371
+ if isinstance(x, FakeTensor):
372
+ input_fake_mode = x.fake_mode
373
+ break
374
+
375
+ if input_fake_mode is not None:
376
+ fake_mode_ctx: Any = patch(
377
+ "torch._inductor.standalone_compile.FakeTensorMode",
378
+ lambda *a, **kw: input_fake_mode,
379
+ )
380
+ else:
381
+ fake_mode_ctx = contextlib.nullcontext()
382
+
383
+ with pregrad_ctx, fake_mode_ctx, _patch_constrain_to_fx_strides():
384
compiled_graph = standalone_compile(graph, example_inputs, **compile_kwargs)
385
386
if use_aot:
vllm/compilation/piecewise_backend.py
+26
-22
@@ -34,13 +34,14 @@ def get_fake_args_from_graph(graph: fx.GraphModule) -> list[Any]:
34
35
36
def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]:
37
- """Create example inputs with symbolic dims replaced by a concrete size.
37
+ """Create Fake example inputs with symbolic dims replaced by a concrete size.
38
39
- Used for single-size eager compilation where we need concrete-shaped
40
- inputs but don't have real runtime tensors yet.
39
+ Used for single-size compilation where we need concrete-shaped inputs.
40
+ The Dynamo-captured graph gives us example inputs with SymInts in them.
41
"""
42
from torch._prims_common import compute_required_storage_length
43
- from torch.fx.experimental.symbolic_shapes import is_symbolic
43
+ from torch._subclasses.fake_tensor import FakeTensorMode
44
+ from torch.fx.experimental.symbolic_shapes import ShapeEnv, is_symbolic
45
46
def concretize(sym_val: Any) -> int:
47
"""Replace all symbolic variables in a SymInt expression with size."""
@@ -49,25 +50,28 @@ def create_concrete_args(graph: fx.GraphModule, size: int) -> list[Any]:
50
expr = sym_val.node.expr
51
return int(expr.subs({s: size for s in expr.free_symbols}))
52
53
+ fake_mode = FakeTensorMode(shape_env=ShapeEnv())
54
+
55
args: list[Any] = []
53
- for node in graph.graph.nodes:
54
- if node.op != "placeholder":
55
- break
56
- val = node.meta["example_value"]
57
- if isinstance(val, torch.SymInt):
58
- args.append(concretize(val))
59
- elif isinstance(val, torch.Tensor):
60
- new_shape = tuple(concretize(d) for d in val.shape)
61
- new_strides = tuple(concretize(s) for s in val.stride())
62
- new_storage_offset = concretize(val.storage_offset())
63
- needed_size = compute_required_storage_length(
64
- new_shape, new_strides, new_storage_offset
65
- )
66
- t = torch.empty(needed_size, dtype=val.dtype, device=val.device)
67
- t = t.as_strided(new_shape, new_strides, new_storage_offset)
68
- args.append(t)
69
- else:
70
- args.append(val)
56
+ with fake_mode:
57
+ for node in graph.graph.nodes:
58
+ if node.op != "placeholder":
59
+ break
60
+ val = node.meta["example_value"]
61
+ if isinstance(val, torch.SymInt):
62
+ args.append(concretize(val))
63
+ elif isinstance(val, torch.Tensor):
64
+ new_shape = tuple(concretize(d) for d in val.shape)
65
+ new_strides = tuple(concretize(s) for s in val.stride())
66
+ new_storage_offset = concretize(val.storage_offset())
67
+ needed_size = compute_required_storage_length(
68
+ new_shape, new_strides, new_storage_offset
69
+ )
70
+ t = torch.empty(needed_size, dtype=val.dtype, device=val.device)
71
+ t = t.as_strided(new_shape, new_strides, new_storage_offset)
72
+ args.append(t)
73
+ else:
74
+ args.append(val)
75
return args
76
77