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 6, 04:30 UTC
Use lazy graph module during split_module to defer recompile() (#37609)
Signed-off-by: angelayi <yiangela7@gmail.com>
Angela Yi committed
Mar 23, 2026 at 08:21 UTC
aceadb5ee1eb0e637706c898b410b18c7b8cc6b2
1 file changed
+8
-3
vllm/compilation/backends.py
+8
-3
@@ -20,6 +20,7 @@ import torch
20
import torch.fx as fx
21
from torch._dynamo.utils import dynamo_timed
22
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.config import CompilationConfig, CUDAGraphMode, VllmConfig
@@ -573,9 +574,13 @@ def split_graph(
574
# otherwise pytorch might reorder the nodes and
575
# the semantics of the graph will change when we
576
# have mutations in the graph
576
- split_gm = torch.fx.passes.split_module.split_module(
577
- graph, None, lambda node: node_to_subgraph_id[node], keep_original_order=True
578
- )
577
+ with _use_lazy_graph_module(True):
578
+ split_gm = torch.fx.passes.split_module.split_module(
579
+ graph,
580
+ None,
581
+ lambda node: node_to_subgraph_id[node],
582
+ keep_original_order=True,
583
+ )
584
585
outputs = []
586