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 4, 06:15 UTC
[Bugfix] Fix codegen for unqualified names (#40726)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
May 6, 2026 at 01:11 UTC
213f10bfdd43422d2936a8f5e218000050572678
4 files changed
+450
-32
tests/compile/test_codegen.py
new
+376
@@ -0,0 +1,376 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Tests for vllm.compilation.codegen — execution code generation.
4
+
5
+Each test runs a real Python function through the same pipeline vLLM uses
6
+in production: ``make_fx`` to obtain an aten-level fx graph, ``split_graph``
7
+to split it into the stitching layer + submodules, and then
8
+``generate_execution_code``/``compile_execution_fn`` for codegen.
9
+"""
10
+
11
+from collections.abc import Callable
12
+
13
+import pytest
14
+import regex as re
15
+import torch
16
+import torch.fx as fx
17
+from torch.fx.experimental.proxy_tensor import make_fx
18
+
19
+from vllm.compilation.backends import split_graph
20
+from vllm.compilation.codegen import (
21
+ _node_ref,
22
+ compile_execution_fn,
23
+ generate_execution_code,
24
+ generate_execution_code_with_name,
25
+)
26
+from vllm.utils.torch_utils import is_torch_equal_or_newer
27
+
28
+
29
+def _trace_and_split(
30
+ model_fn: Callable[..., torch.Tensor],
31
+ example_inputs: tuple[torch.Tensor, ...],
32
+ split_ops: list[str],
33
+) -> fx.GraphModule:
34
+ """Trace ``model_fn`` with make_fx, then split on the named aten ops."""
35
+ gm = make_fx(model_fn)(*example_inputs)
36
+ split_gm, _ = split_graph(gm, split_ops)
37
+ return split_gm
38
+
39
+
40
+def _to_copy_model(x: torch.Tensor) -> torch.Tensor:
41
+ """Traces to ``aten._to_copy.default`` with device + dtype kwargs."""
42
+ return x.to(device=torch.device("cpu"), dtype=torch.float16)
43
+
44
+
45
+def _empty_model(x: torch.Tensor) -> torch.Tensor:
46
+ """Traces to ``aten.empty.memory_format`` with device + dtype kwargs."""
47
+ buf = torch.empty(x.shape, device=torch.device("cpu"), dtype=torch.float16)
48
+ return buf.fill_(0).add(x.to(dtype=torch.float16))
49
+
50
+
51
+@pytest.fixture
52
+def x() -> torch.Tensor:
53
+ return torch.zeros(2, 3)
54
+
55
+
56
+@pytest.mark.parametrize(
57
+ "model_fn,split_ops",
58
+ [
59
+ (_to_copy_model, ["aten::_to_copy.default"]),
60
+ (_empty_model, []),
61
+ ],
62
+ ids=["aten::_to_copy.default", "aten::empty.memory_format"],
63
+)
64
+def test_non_primitive_kwargs_lifted_to_consts(
65
+ model_fn: Callable[[torch.Tensor], torch.Tensor],
66
+ split_ops: list[str],
67
+ x: torch.Tensor,
68
+) -> None:
69
+ """Regression: arguments whose ``repr()`` is not a valid Python
70
+ expression in the generated function's namespace (notably
71
+ ``torch.device``) used to be inlined via ``repr()``, producing source
72
+ like
73
+
74
+ out = torch.ops.aten._to_copy.default(x, device=device(type='cpu'))
75
+
76
+ which fails at call time — only ``torch`` and ``operator`` are imported
77
+ into the namespace, so ``device`` is unbound. The fix collects such
78
+ objects into ``__vllm_consts__`` and references them by index. The
79
+ unqualified ``device(type=...)`` form must never appear in the
80
+ generated source."""
81
+ split_gm = _trace_and_split(model_fn, (x,), split_ops)
82
+ code, submod_names, consts = generate_execution_code(split_gm)
83
+
84
+ assert "device(type=" not in code, (
85
+ "Generated code contains unqualified `device(type=...)` from repr(); "
86
+ "torch.device should be lifted into __vllm_consts__"
87
+ )
88
+ assert torch.device("cpu") in consts, "torch.device kwarg not lifted to consts"
89
+ assert torch.float16 in consts, "torch.dtype kwarg not lifted to consts"
90
+
91
+ fn = compile_execution_fn(code, {}, submod_names, consts)
92
+ out = fn(x)
93
+ expected = model_fn(x)
94
+ assert torch.equal(out, expected), "Compiled output does not match reference"
95
+
96
+
97
+def test_dtype_singleton_deduped(x: torch.Tensor) -> None:
98
+ """``torch.float16`` is a process-wide singleton, so two ops referring
99
+ to it in the traced graph share a single consts slot via ``id()``-based
100
+ dedup. Distinct expressions (``x.to(...)`` vs ``(x*2).to(...)``) ensure
101
+ the tracer can't CSE the two ops into a single node."""
102
+
103
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
104
+ return x.to(dtype=torch.float16) + (x * 2).to(dtype=torch.float16)
105
+
106
+ split_gm = _trace_and_split(model_fn, (x,), [])
107
+ code, submod_names, consts = generate_execution_code(split_gm)
108
+
109
+ # The traced graph must have two distinct _to_copy nodes (otherwise the
110
+ # dedup assertion below is trivially satisfied).
111
+ n_to_copy = sum(
112
+ 1
113
+ for n in split_gm.graph.nodes
114
+ if n.op == "call_module"
115
+ for sn in getattr(split_gm, n.target).graph.nodes
116
+ if sn.op == "call_function" and "to_copy" in sn.name
117
+ )
118
+ assert n_to_copy >= 2, (
119
+ f"Test setup failed: expected ≥2 _to_copy nodes, got {n_to_copy}"
120
+ )
121
+
122
+ assert consts.count(torch.float16) == 1, (
123
+ f"torch.float16 should occupy exactly one slot, got consts={consts}"
124
+ )
125
+ assert code.count("__vllm_consts__[0]") >= 2, (
126
+ "Deduped const slot should be referenced from both _to_copy nodes"
127
+ )
128
+
129
+ fn = compile_execution_fn(code, {}, submod_names, consts)
130
+ assert torch.equal(fn(x), model_fn(x))
131
+
132
+
133
+def test_distinct_dtypes_get_distinct_slots(x: torch.Tensor) -> None:
134
+ """Distinct dtype singletons in the traced graph occupy distinct slots."""
135
+
136
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
137
+ return x.to(dtype=torch.float16) + x.to(dtype=torch.bfloat16)
138
+
139
+ split_gm = _trace_and_split(model_fn, (x,), [])
140
+ _, _, consts = generate_execution_code(split_gm)
141
+
142
+ assert torch.float16 in consts
143
+ assert torch.bfloat16 in consts
144
+ assert len(consts) == 2, f"Expected 2 distinct dtype slots, got {consts}"
145
+
146
+
147
+def test_consts_ordering_deterministic(x: torch.Tensor) -> None:
148
+ """Two independent traces of the same model must produce equal consts
149
+ lists *in the same order*. Cache artifacts identify const slots by
150
+ index, so a non-deterministic order would invalidate cached code."""
151
+
152
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
153
+ # Multiple distinct non-primitives encountered in a fixed graph order.
154
+ a = x.to(device=torch.device("cpu"), dtype=torch.float16)
155
+ return a.to(dtype=torch.bfloat16)
156
+
157
+ _, _, consts1 = generate_execution_code(_trace_and_split(model_fn, (x,), []))
158
+ _, _, consts2 = generate_execution_code(_trace_and_split(model_fn, (x,), []))
159
+
160
+ assert len(consts1) >= 2, "Test setup: model should produce ≥2 const slots"
161
+ assert consts1 == consts2, (
162
+ f"consts ordering must be reproducible across traces; "
163
+ f"got {consts1} vs {consts2}"
164
+ )
165
+
166
+
167
+def test_primitive_args_inlined(x: torch.Tensor) -> None:
168
+ """Primitive args (int dim, etc.) stay inline as repr — no consts."""
169
+
170
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
171
+ return torch.transpose(x, 0, 1).relu()
172
+
173
+ split_gm = _trace_and_split(model_fn, (x,), [])
174
+ code, submod_names, consts = generate_execution_code(split_gm)
175
+
176
+ assert consts == [], "Primitive-only graph must produce empty consts"
177
+
178
+ fn = compile_execution_fn(code, {}, submod_names, consts)
179
+ assert torch.equal(fn(x), model_fn(x))
180
+
181
+
182
+def test_consts_shared_across_split_submods(x: torch.Tensor) -> None:
183
+ """Dedup must apply across inlined submodules, not just within one.
184
+
185
+ The function below splits into three inlined submods, two of which
186
+ independently reference ``torch.float16``. The shared ``const_index``
187
+ threaded through recursive ``generate_execution_code_with_name`` calls
188
+ must collapse the dtype to a single slot used from both submods."""
189
+
190
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
191
+ a = x.to(dtype=torch.float16) # submod_0: _to_copy(fp16)
192
+ b = a.relu() # submod_1: relu (split point)
193
+ c = b.to(dtype=torch.float32) # submod_2: _to_copy(fp32)
194
+ return c.to(dtype=torch.float16) + 1 # submod_2: another _to_copy(fp16)
195
+
196
+ split_gm = _trace_and_split(model_fn, (x,), ["aten::relu.default"])
197
+
198
+ n_submods = sum(1 for _ in split_gm.named_children())
199
+ assert n_submods >= 3, (
200
+ f"Test setup failed: expected ≥3 submods after split, got {n_submods}"
201
+ )
202
+
203
+ code, submod_names, consts = generate_execution_code(split_gm)
204
+
205
+ assert consts.count(torch.float16) == 1, (
206
+ f"fp16 singleton must dedup across submods, got consts={consts}"
207
+ )
208
+
209
+ # Find the consts index for fp16 and confirm at least two distinct
210
+ # inlined submods reference it. This rules out the false-positive where
211
+ # one submod references it twice and the other not at all.
212
+ fp16_idx = consts.index(torch.float16)
213
+ submod_bodies = re.findall(
214
+ r"def __vllm_inlined_submods__(\d+)\([^)]*\):\n((?: .*\n)+)", code
215
+ )
216
+ assert len(submod_bodies) >= 2
217
+ referencing_submods = [
218
+ name for name, body in submod_bodies if f"__vllm_consts__[{fp16_idx}]" in body
219
+ ]
220
+ assert len(referencing_submods) >= 2, (
221
+ f"fp16 slot should be referenced from ≥2 inlined submods, "
222
+ f"got {referencing_submods}"
223
+ )
224
+
225
+ fn = compile_execution_fn(code, {}, submod_names, consts)
226
+ assert torch.equal(fn(x), model_fn(x))
227
+
228
+
229
+def test_non_graphmodule_submod_uses_indexed_callable(x: torch.Tensor) -> None:
230
+ """When a child of split_gm is *not* a ``torch.fx.GraphModule`` — as
231
+ happens in production once ``PiecewiseBackend`` replaces submods —
232
+ codegen emits ``__vllm_submods__[idx](...)`` instead of inlining, and
233
+ the runtime callable is bound from ``submod_callables``."""
234
+
235
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
236
+ return x.relu().sigmoid()
237
+
238
+ split_gm = _trace_and_split(model_fn, (x,), ["aten::relu.default"])
239
+
240
+ # Find a GraphModule child and wrap it in a non-GraphModule nn.Module
241
+ # that delegates to the original — this is the structural shape vLLM
242
+ # produces after PiecewiseBackend takes over a submod.
243
+ child_names = [name for name, _ in split_gm.named_children()]
244
+ target_name = child_names[0]
245
+
246
+ class NonGMWrapper(torch.nn.Module):
247
+ def __init__(self, gm: fx.GraphModule) -> None:
248
+ super().__init__()
249
+ self.gm = gm
250
+
251
+ def forward(self, *args, **kwargs):
252
+ return self.gm(*args, **kwargs)
253
+
254
+ original = getattr(split_gm, target_name)
255
+ del split_gm._modules[target_name]
256
+ split_gm.add_module(target_name, NonGMWrapper(original))
257
+
258
+ code, submod_names, consts = generate_execution_code(split_gm)
259
+
260
+ assert "__vllm_submods__[" in code, (
261
+ "Non-GraphModule submod should produce an indexed callable reference"
262
+ )
263
+ assert target_name in submod_names
264
+
265
+ submod_callables = {
266
+ name: getattr(split_gm, name)
267
+ for name in submod_names
268
+ if not isinstance(getattr(split_gm, name), fx.GraphModule)
269
+ }
270
+ fn = compile_execution_fn(code, submod_callables, submod_names, consts)
271
+ assert torch.equal(fn(x), model_fn(x))
272
+
273
+
274
+# split_graph only passes tuple_return=True to split_module on PyTorch >= 2.12,
275
+# so getitem nodes only appear in the stitching graph from that version onward.
276
+@pytest.mark.skipif(
277
+ not is_torch_equal_or_newer("2.12.0.dev"),
278
+ reason="split_module tuple_return requires PyTorch >= 2.12",
279
+)
280
+def test_getitem_in_stitching_graph(x: torch.Tensor) -> None:
281
+ """``operator.getitem`` on submod tuple returns is the ``call_function``
282
+ special case at codegen.py — emitted as ``name = source[index]``
283
+ rather than a function call."""
284
+
285
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
286
+ return x.relu().sigmoid()
287
+
288
+ split_gm = _trace_and_split(model_fn, (x,), ["aten::relu.default"])
289
+ code, _, _ = generate_execution_code(split_gm)
290
+
291
+ # split_module wraps each submod return in a tuple, so the stitching
292
+ # graph unpacks via getitem. The codegen must emit it as indexing.
293
+ assert re.search(r"\b\w+ = \w+\[\d+\]\n", code), (
294
+ "Stitching graph should emit `name = source[N]` for getitem nodes"
295
+ )
296
+
297
+
298
+def test_del_emitted_for_intermediate_values(x: torch.Tensor) -> None:
299
+ """The codegen schedules ``del`` after a value's last use to free
300
+ memory early. Multi-submod splits naturally have intermediates whose
301
+ last use is not the output node."""
302
+
303
+ def model_fn(x: torch.Tensor) -> torch.Tensor:
304
+ return x.relu().sigmoid().tanh()
305
+
306
+ split_gm = _trace_and_split(
307
+ model_fn, (x,), ["aten::relu.default", "aten::sigmoid.default"]
308
+ )
309
+ code, _, _ = generate_execution_code(split_gm)
310
+
311
+ assert re.search(r"^ del \w+", code, re.MULTILINE), (
312
+ "Liveness analysis should emit `del` for intermediates with "
313
+ "last-use before the output"
314
+ )
315
+
316
+
317
+def test_with_submod_false_rejects_call_module() -> None:
318
+ """``generate_execution_code_with_name(with_submod=False)`` is the
319
+ recursive entry for inlining a GraphModule into its parent. It must
320
+ refuse a graph that itself contains ``call_module`` nodes — the parent
321
+ is responsible for handling those."""
322
+ g = fx.Graph()
323
+ x_node = g.placeholder("x")
324
+ root = torch.nn.Module()
325
+ root.add_module("inner", torch.nn.Identity())
326
+ call = g.call_module("inner", args=(x_node,))
327
+ g.output(call)
328
+ gm = fx.GraphModule(root, g)
329
+
330
+ with pytest.raises(RuntimeError, match="call_module is not allowed"):
331
+ generate_execution_code_with_name(gm, "f", with_submod=False)
332
+
333
+
334
+def test_node_ref_recurses_through_containers() -> None:
335
+ """``_node_ref`` is the recursive walker that lifts non-primitives
336
+ nested inside list/tuple/dict args. Real aten ops rarely produce such
337
+ structures, but the path is needed for DTensor placement lists and
338
+ other future cases — unit-test the walker directly."""
339
+ consts: list = []
340
+ const_index: dict[int, int] = {}
341
+ cpu = torch.device("cpu")
342
+
343
+ # Non-primitive in a list, primitive alongside.
344
+ assert _node_ref([cpu, 1], consts, const_index) == "[__vllm_consts__[0], 1]"
345
+ assert consts == [cpu]
346
+
347
+ # Same object in a tuple — id-based dedup reuses the existing slot.
348
+ assert _node_ref((cpu, 2), consts, const_index) == "(__vllm_consts__[0], 2)"
349
+ assert consts == [cpu]
350
+
351
+ # Single-element tuple uses the trailing-comma form.
352
+ assert _node_ref((cpu,), consts, const_index) == "(__vllm_consts__[0],)"
353
+
354
+ # Dict value lifts the same way.
355
+ ref = _node_ref({"k": cpu}, consts, const_index)
356
+ assert ref == "{'k': __vllm_consts__[0]}"
357
+
358
+
359
+def test_legacy_code_without_consts() -> None:
360
+ """``compile_execution_fn(consts=None)`` must still load code that has
361
+ no ``__vllm_consts__`` reference, so older serialized cache artifacts
362
+ keep working."""
363
+ # Pre-consts codegen: no __vllm_consts__ reference, only torch/operator.
364
+ legacy_code = (
365
+ "import torch\n"
366
+ "def execution_fn(x, *, __vllm_submods__):\n"
367
+ " return __vllm_submods__[0](x) + 1\n"
368
+ )
369
+
370
+ class AddOne(torch.nn.Module):
371
+ def forward(self, x: torch.Tensor) -> torch.Tensor:
372
+ return x + 1
373
+
374
+ fn = compile_execution_fn(legacy_code, {"sub": AddOne()}, ["sub"], consts=None)
375
+ out = fn(torch.zeros(3))
376
+ assert torch.equal(out, torch.full((3,), 2.0))
vllm/compilation/backends.py
+4
-2
@@ -1268,7 +1268,7 @@ class VllmBackend:
1268
original_split_gm if envs.VLLM_USE_MEGA_AOT_ARTIFACT else self.graph
1269
)
1270
1271
- execution_code, submod_names = generate_execution_code(self.split_gm)
1271
+ execution_code, submod_names, consts = generate_execution_code(self.split_gm)
1272
# Use getattr to get correct callables: __dict__ has PiecewiseBackend
1273
# instances (from PiecewiseCompileInterpreter), _modules has originals.
1274
# getattr checks __dict__ first, then falls back to _modules.
@@ -1277,7 +1277,7 @@ class VllmBackend:
1277
for name, _ in self.split_gm.named_children()
1278
}
1279
runtime_callable = compile_execution_fn(
1280
- execution_code, submod_callables, submod_names
1280
+ execution_code, submod_callables, submod_names, consts
1281
)
1282
1283
if (
@@ -1293,6 +1293,7 @@ class VllmBackend:
1293
vllm_backend=self,
1294
execution_code=execution_code,
1295
submod_names=submod_names,
1296
+ consts=consts,
1297
)
1298
1299
# index of tensors that have symbolic shapes (batch size)
@@ -1326,4 +1327,5 @@ class VllmBackend:
1327
sym_tensor_indices=sym_tensor_indices,
1328
execution_code=execution_code,
1329
submod_names=submod_names,
1330
+ consts=consts,
1331
)
vllm/compilation/caching.py
+4
-1
@@ -187,6 +187,7 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
187
aot_autograd_config: dict[str, Any] | None = None,
188
execution_code: str | None = None,
189
submod_names: list[str] | None = None,
190
+ consts: list[Any] | None = None,
191
) -> None:
192
self.graph_module = graph_module
193
self.example_inputs = example_inputs
@@ -198,6 +199,7 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
199
self.sym_tensor_indices = sym_tensor_indices
200
self.execution_code = execution_code
201
self.submod_names = submod_names
202
+ self.consts = consts
203
self._fake_mode: Any | None = None
204
205
import torch._functorch.config as functorch_config
@@ -526,8 +528,9 @@ def reconstruct_serializable_fn_from_mega_artifact(
528
execution_code = state.get("execution_code")
529
submod_names = state.get("submod_names")
530
if execution_code is not None and submod_names is not None:
531
+ consts = state.get("consts")
532
runtime_callable = compile_execution_fn(
530
- execution_code, submod_callables, submod_names
533
+ execution_code, submod_callables, submod_names, consts
534
)
535
else:
536
logger.warning(
vllm/compilation/codegen.py
+66
-29
@@ -22,11 +22,17 @@ def generate_execution_code_with_name(
22
split_gm: torch.fx.GraphModule,
23
fn_name: str,
24
with_submod: bool,
25
-) -> tuple[str, list[str]]:
25
+ consts: list[Any] | None = None,
26
+ const_index: dict[int, int] | None = None,
27
+) -> tuple[str, list[str], list[Any]]:
28
lines: list[str] = []
29
param_names: list[str] = []
30
submod_names: list[str] = []
31
submod_index: dict[str, int] = {}
32
+ if consts is None:
33
+ consts = []
34
+ if const_index is None:
35
+ const_index = {}
36
37
# Build node ordering for liveness analysis.
38
nodes = list(split_gm.graph.nodes)
@@ -48,6 +54,9 @@ def generate_execution_code_with_name(
54
continue
55
del_after.setdefault(node_order[last_user], []).append(node.name)
56
57
+ def ref(arg: Any) -> str:
58
+ return _node_ref(arg, consts, const_index)
59
+
60
for i, node in enumerate(nodes):
61
if node.op == "placeholder":
62
param_names.append(node.name)
@@ -62,16 +71,18 @@ def generate_execution_code_with_name(
71
submod_index[target] = len(submod_names)
72
submod_names.append(target)
73
idx = submod_index[target]
65
- args_str = ", ".join(_node_ref(a) for a in node.args)
66
- kwargs_str = ", ".join(
67
- f"{k}={_node_ref(v)}" for k, v in node.kwargs.items()
68
- )
74
+ args_str = ", ".join(ref(a) for a in node.args)
75
+ kwargs_str = ", ".join(f"{k}={ref(v)}" for k, v in node.kwargs.items())
76
all_args = ", ".join(filter(None, [args_str, kwargs_str]))
77
submod = getattr(split_gm, target)
78
if isinstance(submod, torch.fx.GraphModule):
79
callable_name = f"__vllm_inlined_submods__{idx}"
73
- inlined_code, _ = generate_execution_code_with_name(
74
- submod, callable_name, with_submod=False
80
+ inlined_code, _, _ = generate_execution_code_with_name(
81
+ submod,
82
+ callable_name,
83
+ with_submod=False,
84
+ consts=consts,
85
+ const_index=const_index,
86
)
87
inlined_submods.append(inlined_code)
88
else:
@@ -80,15 +91,13 @@ def generate_execution_code_with_name(
91
92
elif node.op == "call_function":
93
if node.target is operator.getitem:
83
- source = _node_ref(node.args[0])
94
+ source = ref(node.args[0])
95
index = node.args[1]
96
assert isinstance(index, int)
97
lines.append(f" {node.name} = {source}[{index}]")
98
else:
88
- args_str = ", ".join(_node_ref(a) for a in node.args)
89
- kwargs_str = ", ".join(
90
- f"{k}={_node_ref(v)}" for k, v in node.kwargs.items()
91
- )
99
+ args_str = ", ".join(ref(a) for a in node.args)
100
+ kwargs_str = ", ".join(f"{k}={ref(v)}" for k, v in node.kwargs.items())
101
all_args = ", ".join(filter(None, [args_str, kwargs_str]))
102
lines.append(
103
f" {node.name} = {_get_qualified_name(node.target)}({all_args})"
@@ -96,7 +105,7 @@ def generate_execution_code_with_name(
105
106
elif node.op == "output":
107
assert len(node.args) == 1
99
- ret = _node_ref(node.args[0])
108
+ ret = ref(node.args[0])
109
lines.append(f" return {ret}")
110
111
else:
@@ -109,22 +118,29 @@ def generate_execution_code_with_name(
118
119
assert len(param_names) > 0
120
params = ", ".join(param_names)
112
- header = (
113
- f"\ndef {fn_name}({params}{', *, __vllm_submods__' if with_submod else ''}):"
121
+ kw_params = ", *, __vllm_submods__" if with_submod else ""
122
+ header = f"\ndef {fn_name}({params}{kw_params}):"
123
+ return (
124
+ "".join(inlined_submods) + "\n".join([header] + lines) + "\n",
125
+ submod_names,
126
+ consts,
127
)
115
- return "".join(inlined_submods) + "\n".join([header] + lines) + "\n", submod_names
128
129
130
@dynamo_timed("vllm.generate_execution_code")
131
def generate_execution_code(
132
split_gm: torch.fx.GraphModule,
121
-) -> tuple[str, list[str]]:
133
+) -> tuple[str, list[str], list[Any]]:
134
"""Generate Python source code from a split_gm's stitching graph.
135
136
Walks split_gm.graph.nodes and produces a function that calls
137
submodules via a __vllm_submods__ list, avoiding FX GraphModule overhead
138
and dict lookup cost.
139
140
+ Non-primitive constant arguments (e.g. torch.device, DTensor placement
141
+ types) are collected into a constants list and referenced by index
142
+ in the generated code, avoiding reliance on repr() being eval-able.
143
+
144
If a submodule is a plain torch.fx.GraphModule, it is inlined directly
145
in the generated code and we do not need to serialize it in the artifact.
146
@@ -132,15 +148,17 @@ def generate_execution_code(
148
split_gm: The split graph module produced by split_graph().
149
150
Returns:
135
- A tuple of (code, submod_names) where code is the Python source
136
- and submod_names is the ordered list of submodule target names
137
- corresponding to list indices used in the generated code.
151
+ A tuple of (code, submod_names, consts) where code is the Python
152
+ source, submod_names is the ordered list of submodule target names
153
+ corresponding to list indices used in the generated code, and
154
+ consts is a list of non-primitive constant objects referenced
155
+ by the generated code via __vllm_consts__. These objects are
156
+ kept alive for the lifetime of the compiled function.
157
"""
139
-
140
- code, submod_names = generate_execution_code_with_name(
158
+ code, submod_names, consts = generate_execution_code_with_name(
159
split_gm, "execution_fn", with_submod=True
160
)
143
- return "import torch\nimport operator\n" + code, submod_names
161
+ return "import torch\nimport operator\n" + code, submod_names, consts
162
163
164
@dynamo_timed("vllm.compile_execution_fn")
@@ -148,6 +166,7 @@ def compile_execution_fn(
166
code: str,
167
submod_callables: dict[str, Callable[..., Any]],
168
submod_names: list[str],
169
+ consts: list[Any] | None = None,
170
) -> Callable[..., Any]:
171
"""Compile execution code and bind submodule callables.
172
@@ -156,6 +175,9 @@ def compile_execution_fn(
175
submod_callables: Mapping of submodule names to their callables.
176
submod_names: Ordered list of submodule names matching the indices
177
used in the generated code.
178
+ consts: List of non-primitive constant objects referenced by the
179
+ generated code via __vllm_consts__. None for legacy cached
180
+ code that predates this feature.
181
182
Returns:
183
A callable that executes the stitching logic.
@@ -169,6 +191,8 @@ def compile_execution_fn(
191
payload_fn=lambda: code,
192
)
193
namespace: dict[str, Any] = {}
194
+ if consts is not None:
195
+ namespace["__vllm_consts__"] = consts
196
exec(code, namespace) # noqa: S102
197
fn = namespace["execution_fn"]
198
# Using .get() is intentional here because only piecewise backend will
@@ -180,19 +204,32 @@ def compile_execution_fn(
204
return partial(fn, __vllm_submods__=submods_list)
205
206
183
-def _node_ref(arg: Any) -> str:
184
- """Convert an FX node argument to a source code reference recursively."""
207
+def _node_ref(arg: Any, consts: list[Any], const_index: dict[int, int]) -> str:
208
+ """Convert an FX node argument to a source code reference."""
209
if isinstance(arg, torch.fx.Node):
210
return arg.name
211
if isinstance(arg, list):
188
- return f"[{', '.join(_node_ref(x) for x in arg)}]"
212
+ return f"[{', '.join(_node_ref(x, consts, const_index) for x in arg)}]"
213
if isinstance(arg, tuple):
190
- items = ", ".join(_node_ref(x) for x in arg)
214
+ items = ", ".join(_node_ref(x, consts, const_index) for x in arg)
215
return f"({items},)" if len(arg) == 1 else f"({items})"
216
if isinstance(arg, dict):
217
return (
218
"{"
195
- + ", ".join(f"{_node_ref(k)}: {_node_ref(v)}" for k, v in arg.items())
219
+ + ", ".join(
220
+ f"{_node_ref(k, consts, const_index)}: "
221
+ f"{_node_ref(v, consts, const_index)}"
222
+ for k, v in arg.items()
223
+ )
224
+ "}"
225
)
198
- return repr(arg)
226
+ if isinstance(arg, (int, float, bool, str, bytes, type(None))):
227
+ return repr(arg)
228
+ # Dedup by identity, not equality: safe because FX graph args
229
+ # are live for the entire code-generation pass. Objects stored
230
+ # here must be picklable (for compile-artifact caching).
231
+ key = id(arg)
232
+ if key not in const_index:
233
+ const_index[key] = len(consts)
234
+ consts.append(arg)
235
+ return f"__vllm_consts__[{const_index[key]}]"