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 5, 11:00 UTC
[Core] Use tuple_return in split_module for tuple-conformant subgraphs (#38752)
Signed-off-by: Frederik Gossen <frgossen@meta.com> Co-authored-by: Boyuan Feng <boyuan@meta.com>
Frederik Gossen committed
Apr 8, 2026 at 12:09 UTC
d74a306c4b3e7a0b09581af83f82af910a4b758e
1 file changed
+4
vllm/compilation/backends.py
+4
@@ -31,6 +31,7 @@ from vllm.logging_utils import lazy
31
from vllm.platforms import current_platform
32
from vllm.tracing import instrument, instrument_manual
33
from vllm.utils.import_utils import resolve_obj_by_qualname
34
+from vllm.utils.torch_utils import is_torch_equal_or_newer
35
36
from .compiler_interface import (
37
CompilerInterface,
@@ -575,11 +576,14 @@ def split_graph(
576
# the semantics of the graph will change when we
577
# have mutations in the graph
578
with _use_lazy_graph_module(True):
579
+ has_tuple_return = is_torch_equal_or_newer("2.12.0.dev")
580
+ tuple_return_kwarg = {"tuple_return": True} if has_tuple_return else {}
581
split_gm = torch.fx.passes.split_module.split_module(
582
graph,
583
None,
584
lambda node: node_to_subgraph_id[node],
585
keep_original_order=True,
586
+ **tuple_return_kwarg,
587
)
588
589
outputs = []