From aceadb5ee1eb0e637706c898b410b18c7b8cc6b2 Mon Sep 17 00:00:00 2001 From: Angela Yi Date: Mon, 23 Mar 2026 08:21:29 -0700 Subject: [PATCH] Use lazy graph module during split_module to defer recompile() (#37609) Signed-off-by: angelayi --- vllm/compilation/backends.py | 11 ++++++++--- 1 file changed, 8 insertions(+), 3 deletions(-) diff --git a/vllm/compilation/backends.py b/vllm/compilation/backends.py index e049ef345..dee7cdde7 100644 --- a/vllm/compilation/backends.py +++ b/vllm/compilation/backends.py @@ -20,6 +20,7 @@ import torch import torch.fx as fx from torch._dynamo.utils import dynamo_timed from torch._logging._internal import trace_structured +from torch.fx._lazy_graph_module import _use_lazy_graph_module import vllm.envs as envs from vllm.config import CompilationConfig, CUDAGraphMode, VllmConfig @@ -573,9 +574,13 @@ def split_graph( # otherwise pytorch might reorder the nodes and # the semantics of the graph will change when we # have mutations in the graph - split_gm = torch.fx.passes.split_module.split_module( - graph, None, lambda node: node_to_subgraph_id[node], keep_original_order=True - ) + with _use_lazy_graph_module(True): + split_gm = torch.fx.passes.split_module.split_module( + graph, + None, + lambda node: node_to_subgraph_id[node], + keep_original_order=True, + ) outputs = []