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, 18:45 UTC
[compile] Skip FX graph deserialiaztion on loading, further reducing warm compile time. (#40151)
Signed-off-by: zhxchen17 <zhxchen17@fb.com>
Zhengxu Chen committed
Apr 23, 2026 at 01:43 UTC
98a242ff61f7264f6bb7e2d8dd6f78c8496fb39e
3 files changed
+98
-51
vllm/compilation/backends.py
+4
-5
@@ -23,6 +23,10 @@ from torch._logging._internal import trace_structured
23
from torch.fx._lazy_graph_module import _use_lazy_graph_module
24
25
import vllm.envs as envs
26
+from vllm.compilation.codegen import (
27
+ compile_execution_fn,
28
+ generate_execution_code,
29
+)
30
from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig
31
from vllm.config.compilation import DynamicShapesType
32
from vllm.config.utils import Range, hash_factors
@@ -1244,11 +1248,6 @@ class VllmBackend:
1248
original_split_gm if envs.VLLM_USE_MEGA_AOT_ARTIFACT else self.graph
1249
)
1250
1247
- from vllm.compilation.codegen import (
1248
- compile_execution_fn,
1249
- generate_execution_code,
1250
- )
1251
-
1251
execution_code, submod_names = generate_execution_code(self.split_gm)
1252
# Use getattr to get correct callables: __dict__ has PiecewiseBackend
1253
# instances (from PiecewiseCompileInterpreter), _modules has originals.
vllm/compilation/caching.py
+20
-15
@@ -16,6 +16,7 @@ from torch.fx._graph_pickler import GraphPickler, Options
16
from torch.utils import _pytree as pytree
17
18
import vllm.envs as envs
19
+from vllm.compilation.codegen import compile_execution_fn
20
from vllm.compilation.compiler_interface import get_inductor_factors
21
from vllm.compilation.counter import compilation_counter
22
from vllm.config import VllmConfig, get_current_vllm_config
@@ -176,7 +177,7 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
177
178
def __init__(
179
self,
179
- graph_module: torch.fx.GraphModule,
180
+ graph_module: torch.fx.GraphModule | bytes,
181
example_inputs: Sequence[Any],
182
prefix: str,
183
optimized_call: Callable[..., Any],
@@ -187,7 +188,6 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
188
execution_code: str | None = None,
189
submod_names: list[str] | None = None,
190
) -> None:
190
- assert isinstance(graph_module, torch.fx.GraphModule)
191
self.graph_module = graph_module
192
self.example_inputs = example_inputs
193
self.prefix = prefix
@@ -302,10 +302,6 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
302
state = pickle.loads(data)
303
fake_mode = FakeTensorMode(shape_env=ShapeEnv())
304
305
- state["graph_module"] = cls.deserialize_graph_module(
306
- state["graph_module"], fake_mode
307
- )
308
- state["graph_module"].recompile()
305
state["example_inputs"] = GraphPickler.loads(state["example_inputs"], fake_mode)
306
307
standalone_compile_artifacts = state.pop("standalone_compile_artifacts", None)
@@ -331,6 +327,7 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
327
vllm_config=get_current_vllm_config(),
328
sym_shape_indices_map=sym_shape_indices_map,
329
returns_tuple_map=returns_tuple_map,
330
+ fake_mode=fake_mode,
331
)
332
333
logger.info(
@@ -342,6 +339,11 @@ class VllmSerializableFunction(SerializableCallable): # type: ignore[misc]
339
340
return fn
341
342
+ state["graph_module"] = cls.deserialize_graph_module(
343
+ state["graph_module"], fake_mode
344
+ )
345
+ state["graph_module"].recompile()
346
+
347
# Fall back to standard VllmBackend.
348
# Use a lazy closure: the backend needs traced_files for cache
349
# dir computation, but those are only populated after
@@ -410,6 +412,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
412
vllm_config: VllmConfig,
413
sym_shape_indices_map: dict[str, list[int]],
414
returns_tuple_map: dict[str, bool],
415
+ fake_mode: FakeTensorMode,
416
) -> "VllmSerializableFunction":
417
"""Construct a VllmSerializableFunction from cached inductor artifacts.
418
@@ -452,7 +455,6 @@ def reconstruct_serializable_fn_from_mega_artifact(
455
456
prefix = state["prefix"]
457
is_encoder = state.get("is_encoder", False)
455
- split_gm = state["graph_module"]
458
compilation_config = vllm_config.compilation_config
459
460
standalone_compile_artifacts.load_all()
@@ -476,13 +478,16 @@ def reconstruct_serializable_fn_from_mega_artifact(
478
)
479
480
# spot check that cached submodules exist in the graph structure
479
- graph_children = {name for name, _ in split_gm.named_children()}
481
+ # if an old cache is used, this will fail but that's fine because
482
+ # we will just try this error and re-generate the new cache.
483
+ graph_children = set(state["submod_names"])
484
missing = set(piecewise_submod_names) - graph_children
485
assert not missing, (
486
f"artifacts reference submodules not in graph: {missing}. "
487
f"graph has: {sorted(graph_children)}"
488
)
489
490
+ submod_callables = {}
491
for i, submod_name in enumerate(piecewise_submod_names):
492
assert submod_name in sym_shape_indices_map and submod_name in returns_tuple_map
493
@@ -511,7 +516,7 @@ def reconstruct_serializable_fn_from_mega_artifact(
516
is_last,
517
)
518
514
- split_gm.__dict__[submod_name] = wrapped_backend
519
+ submod_callables[submod_name] = wrapped_backend
520
logger.debug(
521
"Replaced submodule %s with piecewise backend from cache",
522
submod_name,
@@ -521,16 +526,16 @@ def reconstruct_serializable_fn_from_mega_artifact(
526
execution_code = state.get("execution_code")
527
submod_names = state.get("submod_names")
528
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
533
+ logger.warning(
534
+ "No execution code found, falling back to graph module execution."
535
+ )
536
+ runtime_callable = GraphPickler.loads(
537
+ state["graph_module"], fake_mode=fake_mode
538
+ )
539
540
if compilation_config.cudagraph_copy_inputs:
541
sym_tensor_indices = state["sym_tensor_indices"]
vllm/compilation/codegen.py
+74
-31
@@ -15,26 +15,14 @@ from typing import Any
15
import torch.fx
16
from torch._dynamo.utils import dynamo_timed
17
from torch._logging import trace_structured
18
+from torch.fx.node import _get_qualified_name
19
20
20
-@dynamo_timed("vllm.generate_execution_code")
21
-def generate_execution_code(
21
+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]]:
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
- """
26
lines: list[str] = []
27
param_names: list[str] = []
28
submod_names: list[str] = []
@@ -43,6 +31,7 @@ def generate_execution_code(
31
# Build node ordering for liveness analysis.
32
nodes = list(split_gm.graph.nodes)
33
node_order = {node: i for i, node in enumerate(nodes)}
34
+ inlined_submods: list[str] = []
35
36
# For each value-producing node, find the position of its last consumer.
37
# If the last consumer is the output node, skip (return handles cleanup).
@@ -65,6 +54,10 @@ def generate_execution_code(
54
55
elif node.op == "call_module":
56
target = node.target
57
+ if not with_submod:
58
+ raise RuntimeError(
59
+ f"call_module is not allowed for codegen target {target}."
60
+ )
61
if target not in submod_index:
62
submod_index[target] = len(submod_names)
63
submod_names.append(target)
@@ -74,13 +67,32 @@ def generate_execution_code(
67
f"{k}={_node_ref(v)}" for k, v in node.kwargs.items()
68
)
69
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}]")
70
+ submod = getattr(split_gm, target)
71
+ if isinstance(submod, torch.fx.GraphModule):
72
+ callable_name = f"__vllm_inlined_submods__{idx}"
73
+ inlined_code, _ = generate_execution_code_with_name(
74
+ submod, callable_name, with_submod=False
75
+ )
76
+ inlined_submods.append(inlined_code)
77
+ else:
78
+ callable_name = f"__vllm_submods__[{idx}]"
79
+ lines.append(f" {node.name} = {callable_name}({all_args})")
80
+
81
+ elif node.op == "call_function":
82
+ if node.target is operator.getitem:
83
+ source = _node_ref(node.args[0])
84
+ index = node.args[1]
85
+ assert isinstance(index, int)
86
+ lines.append(f" {node.name} = {source}[{index}]")
87
+ 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
+ )
92
+ all_args = ", ".join(filter(None, [args_str, kwargs_str]))
93
+ lines.append(
94
+ f" {node.name} = {_get_qualified_name(node.target)}({all_args})"
95
+ )
96
97
elif node.op == "output":
98
assert len(node.args) == 1
@@ -91,14 +103,44 @@ def generate_execution_code(
103
raise RuntimeError(f"Unsupported node from codegen: {node.format_node()}")
104
105
# Emit del for variables whose last use was this node.
94
- if i in del_after:
106
+ if i in del_after and i < len(nodes) - 2:
107
names = sorted(del_after[i])
108
lines.append(f" del {', '.join(names)}")
109
110
assert len(param_names) > 0
111
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
112
+ header = (
113
+ f"\ndef {fn_name}({params}{', *, __vllm_submods__' if with_submod else ''}):"
114
+ )
115
+ return "".join(inlined_submods) + "\n".join([header] + lines) + "\n", submod_names
116
+
117
+
118
+@dynamo_timed("vllm.generate_execution_code")
119
+def generate_execution_code(
120
+ split_gm: torch.fx.GraphModule,
121
+) -> tuple[str, list[str]]:
122
+ """Generate Python source code from a split_gm's stitching graph.
123
+
124
+ Walks split_gm.graph.nodes and produces a function that calls
125
+ submodules via a __vllm_submods__ list, avoiding FX GraphModule overhead
126
+ and dict lookup cost.
127
+
128
+ If a submodule is a plain torch.fx.GraphModule, it is inlined directly
129
+ in the generated code and we do not need to serialize it in the artifact.
130
+
131
+ Args:
132
+ split_gm: The split graph module produced by split_graph().
133
+
134
+ 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.
138
+ """
139
+
140
+ code, submod_names = generate_execution_code_with_name(
141
+ split_gm, "execution_fn", with_submod=True
142
+ )
143
+ return "import torch\nimport operator\n" + code, submod_names
144
145
146
@dynamo_timed("vllm.compile_execution_fn")
@@ -129,11 +171,12 @@ def compile_execution_fn(
171
namespace: dict[str, Any] = {}
172
exec(code, namespace) # noqa: S102
173
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
- ]
174
+ # Using .get() is intentional here because only piecewise backend will
175
+ # be stored in submod_callables. The other submodules are inlined and
176
+ # we don't need to bind them to the execution function. Instead, we
177
+ # should use None as placeholder to ensure the list indices are preserved
178
+ # for better debuggability.
179
+ submods_list = [submod_callables.get(name) for name in submod_names]
180
return partial(fn, __vllm_submods__=submods_list)
181
182