| 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 | from torch.fx.node import _get_qualified_name |
| 19 | |
| 20 | |
| 21 | def generate_execution_code_with_name( |
| 22 | split_gm: torch.fx.GraphModule, |
| 23 | fn_name: str, |
| 24 | with_submod: bool, |
| 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) |
| 39 | node_order = {node: i for i, node in enumerate(nodes)} |
| 40 | inlined_submods: list[str] = [] |
| 41 | |
| 42 | # For each value-producing node, find the position of its last consumer. |
| 43 | # If the last consumer is the output node, skip (return handles cleanup). |
| 44 | # Otherwise, schedule a del after that consumer to free memory early. |
| 45 | del_after: dict[int, list[str]] = {} # position -> names to delete |
| 46 | for node in nodes: |
| 47 | if node.op == "output": |
| 48 | continue |
| 49 | users = list(node.users.keys()) |
| 50 | if not users: |
| 51 | continue |
| 52 | last_user = max(users, key=lambda u: node_order[u]) |
| 53 | if last_user.op == "output": |
| 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) |
| 63 | |
| 64 | elif node.op == "call_module": |
| 65 | target = node.target |
| 66 | if not with_submod: |
| 67 | raise RuntimeError( |
| 68 | f"call_module is not allowed for codegen target {target}." |
| 69 | ) |
| 70 | if target not in submod_index: |
| 71 | submod_index[target] = len(submod_names) |
| 72 | submod_names.append(target) |
| 73 | idx = submod_index[target] |
| 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}" |
| 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: |
| 89 | callable_name = f"__vllm_submods__[{idx}]" |
| 90 | lines.append(f" {node.name} = {callable_name}({all_args})") |
| 91 | |
| 92 | elif node.op == "call_function": |
| 93 | if node.target is operator.getitem: |
| 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: |
| 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})" |
| 104 | ) |
| 105 | |
| 106 | elif node.op == "output": |
| 107 | assert len(node.args) == 1 |
| 108 | ret = ref(node.args[0]) |
| 109 | lines.append(f" return {ret}") |
| 110 | |
| 111 | else: |
| 112 | raise RuntimeError(f"Unsupported node from codegen: {node.format_node()}") |
| 113 | |
| 114 | # Emit del for variables whose last use was this node. |
| 115 | if i in del_after and i < len(nodes) - 2: |
| 116 | names = sorted(del_after[i]) |
| 117 | lines.append(f" del {', '.join(names)}") |
| 118 | |
| 119 | assert len(param_names) > 0 |
| 120 | params = ", ".join(param_names) |
| 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 | ) |
| 128 | |
| 129 | |
| 130 | @dynamo_timed("vllm.generate_execution_code") |
| 131 | def generate_execution_code( |
| 132 | split_gm: torch.fx.GraphModule, |
| 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 | |
| 147 | Args: |
| 148 | split_gm: The split graph module produced by split_graph(). |
| 149 | |
| 150 | Returns: |
| 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 | |
| 158 | """ |
| 159 | code, submod_names, consts = generate_execution_code_with_name( |
| 160 | split_gm, "execution_fn", with_submod=True |
| 161 | ) |
| 162 | return "import torch\nimport operator\n" + code, submod_names, consts |
| 163 | |
| 164 | |
| 165 | @dynamo_timed("vllm.compile_execution_fn") |
| 166 | def compile_execution_fn( |
| 167 | code: str, |
| 168 | submod_callables: dict[str, Callable[..., Any]], |
| 169 | submod_names: list[str], |
| 170 | consts: list[Any] | None = None, |
| 171 | ) -> Callable[..., Any]: |
| 172 | """Compile execution code and bind submodule callables. |
| 173 | |
| 174 | Args: |
| 175 | code: Python source from generate_execution_code(). |
| 176 | submod_callables: Mapping of submodule names to their callables. |
| 177 | submod_names: Ordered list of submodule names matching the indices |
| 178 | used in the generated code. |
| 179 | consts: List of non-primitive constant objects referenced by the |
| 180 | generated code via __vllm_consts__. None for legacy cached |
| 181 | code that predates this feature. |
| 182 | |
| 183 | Returns: |
| 184 | A callable that executes the stitching logic. |
| 185 | |
| 186 | """ |
| 187 | trace_structured( |
| 188 | "artifact", |
| 189 | metadata_fn=lambda: { |
| 190 | "name": "vllm_execution_code", |
| 191 | "encoding": "string", |
| 192 | }, |
| 193 | payload_fn=lambda: code, |
| 194 | ) |
| 195 | namespace: dict[str, Any] = {} |
| 196 | if consts is not None: |
| 197 | namespace["__vllm_consts__"] = consts |
| 198 | exec(code, namespace) # noqa: S102 |
| 199 | fn = namespace["execution_fn"] |
| 200 | # Using .get() is intentional here because only piecewise backend will |
| 201 | # be stored in submod_callables. The other submodules are inlined and |
| 202 | # we don't need to bind them to the execution function. Instead, we |
| 203 | # should use None as placeholder to ensure the list indices are preserved |
| 204 | # for better debuggability. |
| 205 | submods_list = [submod_callables.get(name) for name in submod_names] |
| 206 | return partial(fn, __vllm_submods__=submods_list) |
| 207 | |
| 208 | |
| 209 | def _node_ref(arg: Any, consts: list[Any], const_index: dict[int, int]) -> str: |
| 210 | """Convert an FX node argument to a source code reference.""" |
| 211 | if isinstance(arg, torch.fx.Node): |
| 212 | return arg.name |
| 213 | if isinstance(arg, list): |
| 214 | return f"[{', '.join(_node_ref(x, consts, const_index) for x in arg)}]" |
| 215 | if isinstance(arg, tuple): |
| 216 | items = ", ".join(_node_ref(x, consts, const_index) for x in arg) |
| 217 | return f"({items},)" if len(arg) == 1 else f"({items})" |
| 218 | if isinstance(arg, dict): |
| 219 | return ( |
| 220 | "{" |
| 221 | + ", ".join( |
| 222 | f"{_node_ref(k, consts, const_index)}: " |
| 223 | f"{_node_ref(v, consts, const_index)}" |
| 224 | for k, v in arg.items() |
| 225 | ) |
| 226 | + "}" |
| 227 | ) |
| 228 | if isinstance(arg, (int, float, bool, str, bytes, type(None))): |
| 229 | return repr(arg) |
| 230 | # Dedup by identity, not equality: safe because FX graph args |
| 231 | # are live for the entire code-generation pass. Objects stored |
| 232 | # here must be picklable (for compile-artifact caching). |
| 233 | key = id(arg) |
| 234 | if key not in const_index: |
| 235 | const_index[key] = len(consts) |
| 236 | consts.append(arg) |
| 237 | return f"__vllm_consts__[{const_index[key]}]" |