# SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project """Code generation for split_gm stitching graph execution. Generates a plain Python function that replaces the FX GraphModule's interpreter-based execution of the stitching graph, eliminating nn.Module.__call__ overhead and __getattr__ dispatch. """ import operator from collections.abc import Callable from functools import partial from typing import Any import torch.fx from torch._dynamo.utils import dynamo_timed from torch._logging import trace_structured from torch.fx.node import _get_qualified_name def generate_execution_code_with_name( split_gm: torch.fx.GraphModule, fn_name: str, with_submod: bool, consts: list[Any] | None = None, const_index: dict[int, int] | None = None, ) -> tuple[str, list[str], list[Any]]: lines: list[str] = [] param_names: list[str] = [] submod_names: list[str] = [] submod_index: dict[str, int] = {} if consts is None: consts = [] if const_index is None: const_index = {} # Build node ordering for liveness analysis. nodes = list(split_gm.graph.nodes) node_order = {node: i for i, node in enumerate(nodes)} inlined_submods: list[str] = [] # For each value-producing node, find the position of its last consumer. # If the last consumer is the output node, skip (return handles cleanup). # Otherwise, schedule a del after that consumer to free memory early. del_after: dict[int, list[str]] = {} # position -> names to delete for node in nodes: if node.op == "output": continue users = list(node.users.keys()) if not users: continue last_user = max(users, key=lambda u: node_order[u]) if last_user.op == "output": continue del_after.setdefault(node_order[last_user], []).append(node.name) def ref(arg: Any) -> str: return _node_ref(arg, consts, const_index) for i, node in enumerate(nodes): if node.op == "placeholder": param_names.append(node.name) elif node.op == "call_module": target = node.target if not with_submod: raise RuntimeError( f"call_module is not allowed for codegen target {target}." ) if target not in submod_index: submod_index[target] = len(submod_names) submod_names.append(target) idx = submod_index[target] args_str = ", ".join(ref(a) for a in node.args) kwargs_str = ", ".join(f"{k}={ref(v)}" for k, v in node.kwargs.items()) all_args = ", ".join(filter(None, [args_str, kwargs_str])) submod = getattr(split_gm, target) if isinstance(submod, torch.fx.GraphModule): callable_name = f"__vllm_inlined_submods__{idx}" inlined_code, _, _ = generate_execution_code_with_name( submod, callable_name, with_submod=False, consts=consts, const_index=const_index, ) inlined_submods.append(inlined_code) else: callable_name = f"__vllm_submods__[{idx}]" lines.append(f" {node.name} = {callable_name}({all_args})") elif node.op == "call_function": if node.target is operator.getitem: source = ref(node.args[0]) index = node.args[1] assert isinstance(index, int) lines.append(f" {node.name} = {source}[{index}]") else: args_str = ", ".join(ref(a) for a in node.args) kwargs_str = ", ".join(f"{k}={ref(v)}" for k, v in node.kwargs.items()) all_args = ", ".join(filter(None, [args_str, kwargs_str])) lines.append( f" {node.name} = {_get_qualified_name(node.target)}({all_args})" ) elif node.op == "output": assert len(node.args) == 1 ret = ref(node.args[0]) lines.append(f" return {ret}") else: raise RuntimeError(f"Unsupported node from codegen: {node.format_node()}") # Emit del for variables whose last use was this node. if i in del_after and i < len(nodes) - 2: names = sorted(del_after[i]) lines.append(f" del {', '.join(names)}") assert len(param_names) > 0 params = ", ".join(param_names) kw_params = ", *, __vllm_submods__" if with_submod else "" header = f"\ndef {fn_name}({params}{kw_params}):" return ( "".join(inlined_submods) + "\n".join([header] + lines) + "\n", submod_names, consts, ) @dynamo_timed("vllm.generate_execution_code") def generate_execution_code( split_gm: torch.fx.GraphModule, ) -> tuple[str, list[str], list[Any]]: """Generate Python source code from a split_gm's stitching graph. Walks split_gm.graph.nodes and produces a function that calls submodules via a __vllm_submods__ list, avoiding FX GraphModule overhead and dict lookup cost. Non-primitive constant arguments (e.g. torch.device, DTensor placement types) are collected into a constants list and referenced by index in the generated code, avoiding reliance on repr() being eval-able. If a submodule is a plain torch.fx.GraphModule, it is inlined directly in the generated code and we do not need to serialize it in the artifact. Args: split_gm: The split graph module produced by split_graph(). Returns: A tuple of (code, submod_names, consts) where code is the Python source, submod_names is the ordered list of submodule target names corresponding to list indices used in the generated code, and consts is a list of non-primitive constant objects referenced by the generated code via __vllm_consts__. These objects are kept alive for the lifetime of the compiled function. """ code, submod_names, consts = generate_execution_code_with_name( split_gm, "execution_fn", with_submod=True ) return "import torch\nimport operator\n" + code, submod_names, consts @dynamo_timed("vllm.compile_execution_fn") def compile_execution_fn( code: str, submod_callables: dict[str, Callable[..., Any]], submod_names: list[str], consts: list[Any] | None = None, ) -> Callable[..., Any]: """Compile execution code and bind submodule callables. Args: code: Python source from generate_execution_code(). submod_callables: Mapping of submodule names to their callables. submod_names: Ordered list of submodule names matching the indices used in the generated code. consts: List of non-primitive constant objects referenced by the generated code via __vllm_consts__. None for legacy cached code that predates this feature. Returns: A callable that executes the stitching logic. """ trace_structured( "artifact", metadata_fn=lambda: { "name": "vllm_execution_code", "encoding": "string", }, payload_fn=lambda: code, ) namespace: dict[str, Any] = {} if consts is not None: namespace["__vllm_consts__"] = consts exec(code, namespace) # noqa: S102 fn = namespace["execution_fn"] # Using .get() is intentional here because only piecewise backend will # be stored in submod_callables. The other submodules are inlined and # we don't need to bind them to the execution function. Instead, we # should use None as placeholder to ensure the list indices are preserved # for better debuggability. submods_list = [submod_callables.get(name) for name in submod_names] return partial(fn, __vllm_submods__=submods_list) def _node_ref(arg: Any, consts: list[Any], const_index: dict[int, int]) -> str: """Convert an FX node argument to a source code reference.""" if isinstance(arg, torch.fx.Node): return arg.name if isinstance(arg, list): return f"[{', '.join(_node_ref(x, consts, const_index) for x in arg)}]" if isinstance(arg, tuple): items = ", ".join(_node_ref(x, consts, const_index) for x in arg) return f"({items},)" if len(arg) == 1 else f"({items})" if isinstance(arg, dict): return ( "{" + ", ".join( f"{_node_ref(k, consts, const_index)}: " f"{_node_ref(v, consts, const_index)}" for k, v in arg.items() ) + "}" ) if isinstance(arg, (int, float, bool, str, bytes, type(None))): return repr(arg) # Dedup by identity, not equality: safe because FX graph args # are live for the entire code-generation pass. Objects stored # here must be picklable (for compile-artifact caching). key = id(arg) if key not in const_index: const_index[key] = len(consts) consts.append(arg) return f"__vllm_consts__[{const_index[key]}]"