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 5, 02:15 UTC
[compile] Invoke split FX graph by codegen. (#38657)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Apr 16, 2026 at 00:03 UTC
951dca8019f6cf55f0fd210a004d3f3edfa0d58b
3 files changed
+206
-9
vllm/compilation/backends.py
+23
-2
@@ -1263,6 +1263,23 @@ class VllmBackend:
1263
original_split_gm if envs.VLLM_USE_MEGA_AOT_ARTIFACT else self.graph
1264
)
1265
1266
+ from vllm.compilation.codegen import (
1267
+ compile_execution_fn,
1268
+ generate_execution_code,
1269
+ )
1270
+
1271
+ execution_code, submod_names = 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.
1275
+ submod_callables = {
1276
+ name: getattr(self.split_gm, name)
1277
+ for name, _ in self.split_gm.named_children()
1278
+ }
1279
+ runtime_callable = compile_execution_fn(
1280
+ execution_code, submod_callables, submod_names
1281
+ )
1282
+
1283
if (
1284
self.compilation_config.cudagraph_mode == CUDAGraphMode.NONE
1285
or not self.compilation_config.cudagraph_copy_inputs
@@ -1271,9 +1288,11 @@ class VllmBackend:
1288
graph_to_serialize,
1289
example_inputs,
1290
self.prefix,
1274
- self.split_gm,
1291
+ runtime_callable,
1292
is_encoder=self.is_encoder,
1293
vllm_backend=self,
1294
+ execution_code=execution_code,
1295
+ submod_names=submod_names,
1296
)
1297
1298
# index of tensors that have symbolic shapes (batch size)
@@ -1294,7 +1313,7 @@ class VllmBackend:
1313
copy_and_call = make_copy_and_call(
1314
sym_tensor_indices,
1315
[example_inputs[x].clone() for x in sym_tensor_indices],
1297
- self.split_gm,
1316
+ runtime_callable,
1317
)
1318
1319
return VllmSerializableFunction(
@@ -1305,4 +1324,6 @@ class VllmBackend:
1324
is_encoder=self.is_encoder,
1325
vllm_backend=self,
1326
sym_tensor_indices=sym_tensor_indices,
1327
+ execution_code=execution_code,
1328
+ submod_names=submod_names,
1329
)
vllm/compilation/caching.py
+28
-7
@@ -184,6 +184,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
184
vllm_backend: Any | None = None,
185
sym_tensor_indices: list[int] | None = None,
186
aot_autograd_config: dict[str, Any] | None = None,
187
+ execution_code: str | None = None,
188
+ submod_names: list[str] | None = None,
189
) -> None:
190
assert isinstance(graph_module, torch.fx.GraphModule)
191
self.graph_module = graph_module
@@ -194,6 +196,8 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
196
self.shape_env = None
197
self.vllm_backend = vllm_backend
198
self.sym_tensor_indices = sym_tensor_indices
199
+ self.execution_code = execution_code
200
+ self.submod_names = submod_names
201
self._fake_mode: Any | None = None
202
203
import torch._functorch.config as functorch_config
@@ -453,7 +457,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
457
458
standalone_compile_artifacts.load_all()
459
456
- submod_names = standalone_compile_artifacts.submodule_names()
460
+ piecewise_submod_names = standalone_compile_artifacts.submodule_names()
461
compiled_callables: dict[str, dict[str, Callable[..., Any]]] = {}
462
463
for cache_key in standalone_compile_artifacts.submodule_bytes:
@@ -473,13 +477,13 @@ def reconstruct_serializable_fn_from_mega_artifact(
477
478
# spot check that cached submodules exist in the graph structure
479
graph_children = {name for name, _ in split_gm.named_children()}
476
- missing = set(submod_names) - graph_children
480
+ missing = set(piecewise_submod_names) - graph_children
481
assert not missing, (
482
f"artifacts reference submodules not in graph: {missing}. "
483
f"graph has: {sorted(graph_children)}"
484
)
485
482
- for i, submod_name in enumerate(submod_names):
486
+ for i, submod_name in enumerate(piecewise_submod_names):
487
assert submod_name in sym_shape_indices_map and submod_name in returns_tuple_map
488
489
sym_shape_indices = sym_shape_indices_map[submod_name]
@@ -490,7 +494,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
494
graph=None, # not needed for cached artifacts
495
vllm_config=vllm_config,
496
piecewise_compile_index=i,
493
- total_piecewise_compiles=len(submod_names),
497
+ total_piecewise_compiles=len(piecewise_submod_names),
498
sym_shape_indices=sym_shape_indices,
499
vllm_backend=vllm_backend,
500
returns_tuple=returns_tuple,
@@ -498,7 +502,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
502
)
503
504
is_first = i == 0
501
- is_last = i == len(submod_names) - 1
505
+ is_last = i == len(piecewise_submod_names) - 1
506
wrapped_backend = wrap_with_cudagraph_if_needed(
507
piecewise_backend,
508
vllm_config,
@@ -513,6 +517,21 @@ def reconstruct_serializable_fn_from_mega_artifact(
517
submod_name,
518
)
519
520
+ # Use codegen'd execution code if available, fall back to split_gm
521
+ execution_code = state.get("execution_code")
522
+ submod_names = state.get("submod_names")
523
+ if execution_code is not None and submod_names is not None:
524
+ from vllm.compilation.codegen import compile_execution_fn
525
+
526
+ submod_callables = {
527
+ name: getattr(split_gm, name) for name, _ in split_gm.named_children()
528
+ }
529
+ runtime_callable = compile_execution_fn(
530
+ execution_code, submod_callables, submod_names
531
+ )
532
+ else:
533
+ runtime_callable = split_gm
534
+
535
if compilation_config.cudagraph_copy_inputs:
536
sym_tensor_indices = state["sym_tensor_indices"]
537
input_buffers = [
@@ -521,9 +540,11 @@ def reconstruct_serializable_fn_from_mega_artifact(
540
)
541
for idx in sym_tensor_indices
542
]
524
- optimized_call = make_copy_and_call(sym_tensor_indices, input_buffers, split_gm)
543
+ optimized_call = make_copy_and_call(
544
+ sym_tensor_indices, input_buffers, runtime_callable
545
+ )
546
else:
526
- optimized_call = split_gm
547
+ optimized_call = runtime_callable
548
549
fn = VllmSerializableFunction(
550
**state,
vllm/compilation/codegen.py
new
+155
@@ -0,0 +1,155 @@
1
+# SPDX-License-Identifier: Apache-2.0
2
+# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
+"""Code generation for split_gm stitching graph execution.
4
+
5
+Generates a plain Python function that replaces the FX GraphModule's
6
+interpreter-based execution of the stitching graph, eliminating
7
+nn.Module.__call__ overhead and __getattr__ dispatch.
8
+"""
9
+
10
+import operator
11
+from collections.abc import Callable
12
+from functools import partial
13
+from typing import Any
14
+
15
+import torch.fx
16
+from torch._dynamo.utils import dynamo_timed
17
+from torch._logging import trace_structured
18
+
19
+
20
+@dynamo_timed("vllm.generate_execution_code")
21
+def generate_execution_code(
22
+ split_gm: torch.fx.GraphModule,
23
+) -> tuple[str, list[str]]:
24
+ """Generate Python source code from a split_gm's stitching graph.
25
+
26
+ Walks split_gm.graph.nodes and produces a function that calls
27
+ submodules via a __vllm_submods__ list, avoiding FX GraphModule overhead
28
+ and dict lookup cost.
29
+
30
+ Args:
31
+ split_gm: The split graph module produced by split_graph().
32
+
33
+ Returns:
34
+ A tuple of (code, submod_names) where code is the Python source
35
+ and submod_names is the ordered list of submodule target names
36
+ corresponding to list indices used in the generated code.
37
+ """
38
+ lines: list[str] = []
39
+ param_names: list[str] = []
40
+ submod_names: list[str] = []
41
+ submod_index: dict[str, int] = {}
42
+
43
+ # Build node ordering for liveness analysis.
44
+ nodes = list(split_gm.graph.nodes)
45
+ node_order = {node: i for i, node in enumerate(nodes)}
46
+
47
+ # For each value-producing node, find the position of its last consumer.
48
+ # If the last consumer is the output node, skip (return handles cleanup).
49
+ # Otherwise, schedule a del after that consumer to free memory early.
50
+ del_after: dict[int, list[str]] = {} # position -> names to delete
51
+ for node in nodes:
52
+ if node.op == "output":
53
+ continue
54
+ users = list(node.users.keys())
55
+ if not users:
56
+ continue
57
+ last_user = max(users, key=lambda u: node_order[u])
58
+ if last_user.op == "output":
59
+ continue
60
+ del_after.setdefault(node_order[last_user], []).append(node.name)
61
+
62
+ for i, node in enumerate(nodes):
63
+ if node.op == "placeholder":
64
+ param_names.append(node.name)
65
+
66
+ elif node.op == "call_module":
67
+ target = node.target
68
+ if target not in submod_index:
69
+ submod_index[target] = len(submod_names)
70
+ submod_names.append(target)
71
+ idx = submod_index[target]
72
+ args_str = ", ".join(_node_ref(a) for a in node.args)
73
+ kwargs_str = ", ".join(
74
+ f"{k}={_node_ref(v)}" for k, v in node.kwargs.items()
75
+ )
76
+ all_args = ", ".join(filter(None, [args_str, kwargs_str]))
77
+ lines.append(f" {node.name} = __vllm_submods__[{idx}]({all_args})")
78
+
79
+ elif node.op == "call_function" and node.target is operator.getitem:
80
+ source = _node_ref(node.args[0])
81
+ index = node.args[1]
82
+ assert isinstance(index, int)
83
+ lines.append(f" {node.name} = {source}[{index}]")
84
+
85
+ elif node.op == "output":
86
+ assert len(node.args) == 1
87
+ ret = _node_ref(node.args[0])
88
+ lines.append(f" return {ret}")
89
+
90
+ else:
91
+ raise RuntimeError(f"Unsupported node from codegen: {node.format_node()}")
92
+
93
+ # Emit del for variables whose last use was this node.
94
+ if i in del_after:
95
+ names = sorted(del_after[i])
96
+ lines.append(f" del {', '.join(names)}")
97
+
98
+ assert len(param_names) > 0
99
+ params = ", ".join(param_names)
100
+ header = f"def execution_fn({params}, *, __vllm_submods__):"
101
+ return "import torch\n" + "\n".join([header] + lines) + "\n", submod_names
102
+
103
+
104
+@dynamo_timed("vllm.compile_execution_fn")
105
+def compile_execution_fn(
106
+ code: str,
107
+ submod_callables: dict[str, Callable[..., Any]],
108
+ submod_names: list[str],
109
+) -> Callable[..., Any]:
110
+ """Compile execution code and bind submodule callables.
111
+
112
+ Args:
113
+ code: Python source from generate_execution_code().
114
+ submod_callables: Mapping of submodule names to their callables.
115
+ submod_names: Ordered list of submodule names matching the indices
116
+ used in the generated code.
117
+
118
+ Returns:
119
+ A callable that executes the stitching logic.
120
+ """
121
+ trace_structured(
122
+ "artifact",
123
+ metadata_fn=lambda: {
124
+ "name": "vllm_execution_code",
125
+ "encoding": "string",
126
+ },
127
+ payload_fn=lambda: code,
128
+ )
129
+ namespace: dict[str, Any] = {}
130
+ exec(code, namespace) # noqa: S102
131
+ fn = namespace["execution_fn"]
132
+ # Use .forward() directly to avoid nn.Module.__call__ overhead.
133
+ submods_list = [
134
+ c.forward if isinstance(c, torch.fx.GraphModule) else c
135
+ for c in (submod_callables[name] for name in submod_names)
136
+ ]
137
+ return partial(fn, __vllm_submods__=submods_list)
138
+
139
+
140
+def _node_ref(arg: Any) -> str:
141
+ """Convert an FX node argument to a source code reference recursively."""
142
+ if isinstance(arg, torch.fx.Node):
143
+ return arg.name
144
+ if isinstance(arg, list):
145
+ return f"[{', '.join(_node_ref(x) for x in arg)}]"
146
+ if isinstance(arg, tuple):
147
+ items = ", ".join(_node_ref(x) for x in arg)
148
+ return f"({items},)" if len(arg) == 1 else f"({items})"
149
+ if isinstance(arg, dict):
150
+ return (
151
+ "{"
152
+ + ", ".join(f"{_node_ref(k)}: {_node_ref(v)}" for k, v in arg.items())
153
+ + "}"
154
+ )
155
+ return repr(arg)