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 10, 09:00 UTC
[Misc][LLaMa4] Compile LLaMa Vision Encoder (#30709)
Signed-off-by: Lucas Kabela <lucaskabela@meta.com>
Lucas Kabela committed
Jan 9, 2026 at 19:01 UTC
ea6d067a2aebc32976f4ebfc0ab2546fbaa91c42
7 files changed
+85
-20
tests/compile/fullgraph/test_multimodal_compile.py
+37
@@ -71,3 +71,40 @@ def test_qwen2_5_vl_no_vit_compilation(vllm_runner, monkeypatch):
71
) as _,
72
):
73
pass
74
+
75
+
76
+# forked needed to workaround https://github.com/vllm-project/vllm/issues/21073
77
+# Requires Cuda and 8 gpus as well
78
+@pytest.mark.forked
79
+@pytest.mark.skip(reason="Skipping due to CI resource constraints")
80
+def test_mllama4_vit_compilation(vllm_runner, monkeypatch):
81
+ """Test that Mllama4 vision submodules are compiled.
82
+
83
+ This test verifies that the 2 vision submodules (Llama4VisionEncoder,
84
+ Llama4VisionPixelShuffleMLP) are properly tagged
85
+ for compilation by checking that num_models_seen increases to 3.
86
+
87
+ However since we are using TP=8, we compilation_counter will not
88
+ work properly so we will just check the run succeeds rn
89
+ """
90
+ # Disable multiprocessing so that the counter is in the same process
91
+ monkeypatch.setenv("VLLM_ENABLE_V1_MULTIPROCESSING", "0")
92
+
93
+ with (
94
+ monkeypatch.context(),
95
+ # TODO: Since we require TP=8, this messes with the compilation
96
+ # counter. We should fix this in the future, but leave for now
97
+ # to make sure that compilation runs (no crash) with llama vision encoder
98
+ compilation_counter.expect(num_models_seen=0),
99
+ vllm_runner(
100
+ "meta-llama/Llama-4-Scout-17B-16E-Instruct",
101
+ max_model_len=512,
102
+ gpu_memory_utilization=0.8,
103
+ tensor_parallel_size=8,
104
+ compilation_config={
105
+ "mode": CompilationMode.VLLM_COMPILE,
106
+ "compile_mm_encoder": True,
107
+ },
108
+ ),
109
+ ):
110
+ pass
vllm/config/compilation.py
+3
-2
@@ -430,8 +430,9 @@ class CompilationConfig:
430
If empty list [], no ops are excluded (suitable for full cudagraphs)."""
431
compile_mm_encoder: bool = False
432
"""Whether or not to compile the multimodal encoder.
433
- Currently, this only works for `Qwen2_5_vl` on selected platforms.
434
- Disabled by default until more models are supported/tested to work."""
433
+ Currently, this only works for `Qwen2_5_vl` and `mLLaMa4` models
434
+ on selected platforms. Disabled by default until more models
435
+ are supported/tested to work."""
436
437
# Inductor capture
438
compile_sizes: list[int | str] | None = None
vllm/model_executor/layers/attention/mm_encoder_attention.py
+3
-3
@@ -171,12 +171,12 @@ class MMEncoderAttention(CustomOp):
171
q=query,
172
k=key,
173
v=value,
174
- scale=self.scale,
175
- cu_seqlens=cu_seqlens,
176
- max_seqlen=max_seqlen,
174
batch_size=bsz,
175
is_rocm_aiter=(self.attn_backend == AttentionBackendEnum.ROCM_AITER_FA),
176
fa_version=self._fa_version,
177
+ scale=self.scale,
178
+ cu_seqlens=cu_seqlens,
179
+ max_seqlen=max_seqlen,
180
)
181
if is_reshaped:
182
output = output.reshape(bsz, q_len, -1)
vllm/model_executor/layers/rotary_embedding/llama4_vision_rope.py
+5
-2
@@ -60,14 +60,17 @@ class Llama4VisionRotaryEmbedding(RotaryEmbeddingBase):
60
assert key is not None
61
# self.cos_sin_cache here is complex tensor so we cannot cast into
62
# query's dtype directly with self._match_cos_sin_cache_dtype
63
- self.cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)
63
+
64
+ # NOTE: by not storing cos_sin_cache in self, we can avoid
65
+ # memory buffer update which is costly to runtime
66
+ cos_sin_cache: torch.Tensor = self.cos_sin_cache.to(query.device)
67
query_ = torch.view_as_complex(query.float().reshape(*query.shape[:-1], -1, 2))
68
key_ = torch.view_as_complex(key.float().reshape(*key.shape[:-1], -1, 2))
69
broadcast_shape = [
70
d if i == 1 or i == (query_.ndim - 1) else 1
71
for i, d in enumerate(query_.shape)
72
]
70
- freqs_ci = self.cos_sin_cache.view(*broadcast_shape)
73
+ freqs_ci = cos_sin_cache.view(*broadcast_shape)
74
query_out = torch.view_as_real(query_ * freqs_ci).flatten(3)
75
key_out = torch.view_as_real(key_ * freqs_ci).flatten(3)
76
return query_out.type_as(query), key_out.type_as(key)
vllm/model_executor/models/llama.py
+5
-1
@@ -369,7 +369,11 @@ def llama_model_invariants(
369
torch._check(positions.size()[0] == input_ids.size()[0])
370
371
372
-@support_torch_compile(shape_invariants=llama_model_invariants)
372
+@support_torch_compile(
373
+ # TODO[#32068]: Investigate recompilation
374
+ # mark_unbacked_dims={"input_ids": 0},
375
+ shape_invariants=llama_model_invariants
376
+)
377
class LlamaModel(nn.Module):
378
def __init__(
379
self,
vllm/model_executor/models/mllama4.py
+29
-9
@@ -31,9 +31,11 @@ from transformers.models.llama4.image_processing_llama4_fast import (
31
get_best_fit,
32
)
33
34
-from vllm.config import VllmConfig
34
+from vllm.compilation.decorators import support_torch_compile
35
+from vllm.config import VllmConfig, set_current_vllm_config
36
from vllm.config.multimodal import BaseDummyOptions
37
from vllm.distributed import get_tensor_model_parallel_world_size
38
+from vllm.forward_context import set_forward_context
39
from vllm.model_executor.layers.attention.mm_encoder_attention import MMEncoderAttention
40
from vllm.model_executor.layers.fused_moe import FusedMoE
41
from vllm.model_executor.layers.linear import (
@@ -47,6 +49,7 @@ from vllm.model_executor.layers.rotary_embedding import get_rope
49
from vllm.model_executor.model_loader.utils import initialize_model
50
from vllm.model_executor.model_loader.weight_utils import default_weight_loader
51
from vllm.model_executor.models.module_mapping import MultiModelKeys
52
+from vllm.model_executor.models.vision import should_torch_compile_mm_vit
53
from vllm.multimodal import MULTIMODAL_REGISTRY
54
from vllm.multimodal.inputs import (
55
MultiModalDataDict,
@@ -456,6 +459,9 @@ class Llama4UnfoldConvolution(nn.Module):
459
return hidden_states
460
461
462
+@support_torch_compile(
463
+ dynamic_arg_dims={"images_flattened": 0}, enable_if=should_torch_compile_mm_vit
464
+)
465
class Llama4VisionModel(nn.Module):
466
def __init__(
467
self,
@@ -497,6 +503,7 @@ class Llama4VisionModel(nn.Module):
503
prefix=f"{prefix}.model",
504
use_data_parallel=use_data_parallel,
505
)
506
+
507
self.vision_adapter = Llama4VisionPixelShuffleMLP(
508
config,
509
quant_config,
@@ -762,18 +769,28 @@ class Llama4ForConditionalGeneration(
769
multimodal_config = vllm_config.model_config.multimodal_config
770
self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"
771
772
+ self.vllm_config = vllm_config
773
self.config = config
774
self.quant_config = quant_config
775
self.multimodal_config = multimodal_config
776
if multimodal_config.get_limit_per_prompt("image"):
769
- self.vision_model = Llama4VisionModel(
770
- config.vision_config,
771
- None,
772
- prefix=maybe_prefix(prefix, "vision_model"),
773
- use_data_parallel=self.use_data_parallel,
774
- )
777
+ from vllm.compilation.backends import set_model_tag
778
+
779
+ with (
780
+ set_current_vllm_config(vllm_config),
781
+ set_model_tag("Llama4VisionModel", is_encoder=True),
782
+ ):
783
+ self.vision_model = Llama4VisionModel(
784
+ config=config.vision_config,
785
+ quant_config=None,
786
+ prefix=maybe_prefix(prefix, "vision_model"),
787
+ use_data_parallel=self.use_data_parallel,
788
+ )
789
+
790
self.multi_modal_projector = Llama4MultiModalProjector(
776
- self.config, None, prefix=maybe_prefix(prefix, "multi_modal_projector")
791
+ config=self.config,
792
+ quant_config=None,
793
+ prefix=maybe_prefix(prefix, "multi_modal_projector"),
794
)
795
else:
796
self.vision_model = None
@@ -883,7 +900,10 @@ class Llama4ForConditionalGeneration(
900
if image_input is None:
901
return []
902
886
- return self._process_image_input(image_input)
903
+ with (
904
+ set_forward_context(None, self.vllm_config),
905
+ ):
906
+ return self._process_image_input(image_input)
907
908
def forward(
909
self,
vllm/v1/attention/ops/vit_attn_wrappers.py
+3
-3
@@ -72,9 +72,9 @@ def flash_attn_maxseqlen_wrapper_fake(
72
batch_size: int,
73
is_rocm_aiter: bool,
74
fa_version: int | None,
75
- scale: float | None,
76
- cu_seqlens: torch.Tensor | None,
77
- max_seqlen: torch.Tensor | None,
75
+ scale: float | None = None,
76
+ cu_seqlens: torch.Tensor | None = None,
77
+ max_seqlen: torch.Tensor | None = None,
78
) -> torch.Tensor:
79
return torch.empty_like(q)
80