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 11, 12:30 UTC
[torch.compile] Add encoder tag for compilation (#30489)
Signed-off-by: ilmarkov <markovilya197@gmail.com>
Ilya Markov committed
Dec 14, 2025 at 11:15 UTC
3224ea9915750cdd714d85c843264923ef4018cc
3 files changed
+14
-10
vllm/compilation/backends.py
+10
-1
@@ -463,21 +463,27 @@ class PiecewiseCompileInterpreter(torch.fx.Interpreter):
463
# the tag for the part of model being compiled,
464
# e.g. backbone/eagle_head
465
model_tag: str = "backbone"
466
+model_is_encoder: bool = False
467
468
469
@contextmanager
469
-def set_model_tag(tag: str):
470
+def set_model_tag(tag: str, is_encoder: bool = False):
471
"""Context manager to set the model tag."""
472
global model_tag
473
+ global model_is_encoder
474
assert tag != model_tag, (
475
f"Model tag {tag} is the same as the current tag {model_tag}."
476
)
477
old_tag = model_tag
478
+ old_is_encoder = model_is_encoder
479
+
480
model_tag = tag
481
+ model_is_encoder = is_encoder
482
try:
483
yield
484
finally:
485
model_tag = old_tag
486
+ model_is_encoder = old_is_encoder
487
488
489
class VllmBackend:
@@ -523,6 +529,9 @@ class VllmBackend:
529
# them, e.g. backbone (default), eagle_head, etc.
530
self.prefix = prefix or model_tag
531
532
+ # Mark compilation for encoder.
533
+ self.is_encoder = model_is_encoder
534
+
535
# Passes to run on the graph post-grad.
536
self.pass_manager = resolve_obj_by_qualname(
537
current_platform.get_pass_manager_cls()
vllm/compilation/piecewise_backend.py
+1
-6
@@ -53,12 +53,7 @@ class PiecewiseBackend:
53
self.is_last_graph = piecewise_compile_index == total_piecewise_compiles - 1
54
55
self.is_full_graph = total_piecewise_compiles == 1
56
- # TODO: we need to generalize encoder compilation to other models
57
- self.is_encoder_compilation = vllm_backend.prefix in [
58
- "Qwen2_5_VisionPatchEmbed",
59
- "Qwen2_5_VisionPatchMerger",
60
- "Qwen2_5_VisionBlock",
61
- ]
56
+ self.is_encoder_compilation = vllm_backend.is_encoder
57
58
self.compile_ranges = self.compilation_config.get_compile_ranges()
59
if self.is_encoder_compilation:
vllm/model_executor/models/qwen2_5_vl.py
+3
-3
@@ -612,7 +612,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
612
# DO NOT MOVE THIS IMPORT
613
from vllm.compilation.backends import set_model_tag
614
615
- with set_model_tag("Qwen2_5_VisionPatchEmbed"):
615
+ with set_model_tag("Qwen2_5_VisionPatchEmbed", is_encoder=True):
616
self.patch_embed = Qwen2_5_VisionPatchEmbed(
617
patch_size=patch_size,
618
temporal_patch_size=temporal_patch_size,
@@ -651,7 +651,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
651
f"Qwen2.5-VL does not support {self.attn_backend} backend now."
652
)
653
654
- with set_model_tag("Qwen2_5_VisionBlock"):
654
+ with set_model_tag("Qwen2_5_VisionBlock", is_encoder=True):
655
self.blocks = nn.ModuleList(
656
[
657
Qwen2_5_VisionBlock(
@@ -670,7 +670,7 @@ class Qwen2_5_VisionTransformer(nn.Module):
670
]
671
)
672
673
- with set_model_tag("Qwen2_5_VisionPatchMerger"):
673
+ with set_model_tag("Qwen2_5_VisionPatchMerger", is_encoder=True):
674
self.merger = Qwen2_5_VisionPatchMerger(
675
d_model=vision_config.out_hidden_size,
676
context_dim=self.hidden_size,