mllama4.py 41.7 KB
Newer Older
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
#
# Copyright 2025 the LLAMA4, Meta Inc., vLLM, and HuggingFace Inc. team.
# All rights reserved.
#
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
import math
from collections.abc import Iterable, Mapping
from itertools import tee
22
from typing import Annotated, Literal
23
24
25
26
27
28
29

import torch
from torch import nn
from transformers import BatchFeature, Llama4Config, Llama4VisionConfig
from transformers.image_utils import SizeDict
from transformers.models.llama4 import Llama4Processor
from transformers.models.llama4.image_processing_llama4_fast import (
30
31
32
    find_supported_resolutions,
    get_best_fit,
)
33

34
35
from vllm.compilation.decorators import support_torch_compile
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
42
43
44
45
46
from vllm.model_executor.layers.linear import (
    ColumnParallelLinear,
    QKVParallelLinear,
    ReplicatedLinear,
    RowParallelLinear,
)
47
48
from vllm.model_executor.layers.quantization import QuantizationConfig
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
55
56
57
58
59
60
61
62
63
64
65
66
67
from vllm.multimodal.inputs import (
    MultiModalDataDict,
    MultiModalFieldConfig,
    MultiModalKwargsItems,
)
from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
from vllm.multimodal.processing import (
    BaseMultiModalProcessor,
    BaseProcessingInfo,
    InputProcessingContext,
    PromptReplacement,
    PromptUpdate,
    PromptUpdateDetails,
)
68
from vllm.multimodal.profiling import BaseDummyInputsBuilder
69
from vllm.sequence import IntermediateTensors
70
from vllm.utils.tensor_schema import TensorSchema, TensorShape
71

72
from .interfaces import (
73
    MixtureOfExperts,
74
75
    MultiModalEmbeddings,
    SupportsEagle3,
76
    SupportsLoRA,
77
78
79
    SupportsMultiModal,
    SupportsPP,
)
80
from .llama4 import Llama4ForCausalLM
81
82
83
84
from .utils import (
    AutoWeightsLoader,
    maybe_prefix,
)
85
from .vision import run_dp_sharded_vision_model
86
87


88
class Llama4ImagePatchInputs(TensorSchema):
89
    """
90
91
92
93
94
    Dimensions:
        - batch_size: Batch size
        - total_num_chunks: Batch size * number of chunks
        - num_channels: Number of channels
        - image_size: Size of each image
95
    """
96
97
98

    type: Literal["pixel_values"] = "pixel_values"

99
    pixel_values: Annotated[
100
101
102
        torch.Tensor,
        TensorShape("total_num_chunks", "num_channels", "image_size", "image_size"),
    ]
103
104

    patches_per_image: Annotated[torch.Tensor, TensorShape("batch_size")]
105
106
    """
    The number of total patches for each image in the batch.
107
    
108
    This is used to split the embeddings which has the first two dimensions
109
    flattened just like `pixel_values`.
110
    """
111

112
    aspect_ratios: Annotated[torch.Tensor, TensorShape("batch_size", 2)]
113
114
115
    """
    A list of aspect ratios corresponding to the number of tiles
    in each dimension that each image in the batch corresponds to.
116
    Each aspect ratio is a pair (ratio_h, ratio_w).
117
118
119
120
    """


class Llama4VisionMLP(nn.Module):
121
122
123
124
125
126
127
    def __init__(
        self,
        input_size: int,
        intermediate_size: int,
        output_size: int,
        bias: bool,
        output_activation: bool,
128
        quant_config: QuantizationConfig | None = None,
129
130
131
        prefix: str = "",
        use_data_parallel: bool = False,
    ):
132
        super().__init__()
133
        self.fc1 = ColumnParallelLinear(
134
135
136
137
138
            input_size=input_size,
            output_size=intermediate_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.fc1",
139
            disable_tp=use_data_parallel,
140
        )
141
        self.fc2 = RowParallelLinear(
142
143
144
145
146
            input_size=intermediate_size,
            output_size=output_size,
            bias=bias,
            quant_config=quant_config,
            prefix=f"{prefix}.fc2",
147
            disable_tp=use_data_parallel,
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
        )
        self.activation_fn = nn.GELU()
        self.output_activation = output_activation

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states, _ = self.fc1(hidden_states)
        hidden_states = self.activation_fn(hidden_states)
        hidden_states, _ = self.fc2(hidden_states)
        if self.output_activation:
            return self.activation_fn(hidden_states)
        return hidden_states


class Llama4MultiModalProjector(nn.Module):
    def __init__(
        self,
        config,
165
        quant_config: QuantizationConfig | None = None,
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
        prefix: str = "",
    ):
        super().__init__()
        self.linear_1 = ColumnParallelLinear(
            input_size=config.vision_config.vision_output_dim,
            output_size=config.text_config.hidden_size,
            bias=False,
            quant_config=quant_config,
            gather_output=True,
            prefix=f"{prefix}.linear_1",
        )

    def forward(self, image_features):
        hidden_states, _ = self.linear_1(image_features)
        return hidden_states


def pixel_shuffle(input_tensor, shuffle_ratio):
    # input_tensor: [batch_size, num_patches, channels]
    batch_size, num_patches, channels = input_tensor.shape
    patch_size = int(math.sqrt(num_patches))

    input_tensor = input_tensor.view(batch_size, patch_size, patch_size, -1)
    batch_size, height, width, channels = input_tensor.size()

191
192
193
    reshaped_tensor = input_tensor.view(
        batch_size, height, int(width * shuffle_ratio), int(channels / shuffle_ratio)
    )
194
195
    reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()

196
197
198
199
200
201
    reshaped_tensor = reshaped_tensor.view(
        batch_size,
        int(height * shuffle_ratio),
        int(width * shuffle_ratio),
        int(channels / (shuffle_ratio**2)),
    )
202
203
    reshaped_tensor = reshaped_tensor.permute(0, 2, 1, 3).contiguous()

204
    output_tensor = reshaped_tensor.view(batch_size, -1, reshaped_tensor.shape[-1])
205
206
207
208
209
210
211
    return output_tensor


class Llama4VisionPixelShuffleMLP(nn.Module):
    def __init__(
        self,
        config,
212
        quant_config: QuantizationConfig | None = None,
213
        prefix: str = "",
214
        use_data_parallel: bool = False,
215
216
217
    ):
        super().__init__()
        self.pixel_shuffle_ratio = config.pixel_shuffle_ratio
218
219
220
        self.inner_dim = int(
            config.projector_input_dim // (self.pixel_shuffle_ratio**2)
        )
221
222
223
224
225
226
227
228
        self.output_dim = config.projector_output_dim
        self.mlp = Llama4VisionMLP(
            input_size=config.intermediate_size,
            intermediate_size=config.projector_input_dim,
            output_size=config.projector_output_dim,
            bias=config.multi_modal_projector_bias,
            output_activation=True,
            quant_config=quant_config,
229
230
231
            prefix=f"{prefix}.mlp",
            use_data_parallel=use_data_parallel,
        )
232
233

    def forward(self, encoded_patches: torch.Tensor) -> torch.Tensor:
234
        encoded_patches = pixel_shuffle(encoded_patches, self.pixel_shuffle_ratio)
235
236
237
238
239
240
241
        return self.mlp(encoded_patches)


class Llama4VisionAttention(nn.Module):
    def __init__(
        self,
        config: Llama4VisionConfig,
242
        quant_config: QuantizationConfig | None,
243
        prefix: str = "",
244
        use_data_parallel: bool = False,
245
246
247
    ):
        super().__init__()
        self.config = config
248
249
250
        self.tp_size = (
            1 if use_data_parallel else get_tensor_model_parallel_world_size()
        )
251
252
253
254
255
256
257
258
259
260
        self.embed_dim = config.hidden_size
        self.num_heads = config.num_attention_heads
        self.head_dim = config.hidden_size // self.num_heads
        assert self.num_heads % self.tp_size == 0
        self.num_local_heads = self.num_heads // self.tp_size
        self.q_size = self.num_local_heads * self.head_dim
        self.kv_size = self.num_local_heads * self.head_dim
        self.attention_dropout = config.attention_dropout
        self.scaling = self.head_dim**-0.5

261
        self.attn = MMEncoderAttention(
262
263
            self.num_local_heads, self.head_dim, self.scaling
        )
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296

        if use_data_parallel:
            self.qkv_proj = ReplicatedLinear(
                self.embed_dim,
                self.q_size + 2 * self.kv_size,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.qkv_proj",
            )
            self.o_proj = ReplicatedLinear(
                self.num_heads * self.head_dim,
                self.embed_dim,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.o_proj",
            )
        else:
            self.qkv_proj = QKVParallelLinear(
                self.embed_dim,
                self.head_dim,
                self.num_heads,
                bias=True,
                quant_config=quant_config,
                prefix=f"{prefix}.qkv_proj",
            )
            self.o_proj = RowParallelLinear(
                self.num_heads * self.head_dim,
                self.embed_dim,
                bias=True,
                input_is_parallel=True,
                quant_config=quant_config,
                prefix=f"{prefix}.o_proj",
            )
297

298
299
300
        rope_parameters = {
            "rope_type": "mllama4",
            "rope_theta": config.rope_parameters["rope_theta"],
301
            "partial_rotary_factor": 0.5,
302
303
        }

304
305
306
        self.rotary_emb = get_rope(
            head_size=self.head_dim,
            # number of image patches
307
            max_position=(config.image_size // config.patch_size) ** 2,
308
            rope_parameters=rope_parameters,
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
            is_neox_style=False,
            dtype=torch.complex64,  # important
        )

    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        input_shape = hidden_states.shape[:-1]

        qkv, _ = self.qkv_proj(hidden_states)
        q, k, v = qkv.split([self.q_size, self.kv_size, self.kv_size], dim=-1)

        q = q.view(q.shape[0], q.shape[1], self.num_local_heads, self.head_dim)
        k = k.view(k.shape[0], k.shape[1], self.num_local_heads, self.head_dim)
        q, k = self.rotary_emb(q, k)

        q = q.view(q.shape[0], q.shape[1], -1)
        k = k.view(k.shape[0], k.shape[1], -1)

        attn_output = self.attn(q, k, v)
        attn_output = attn_output.reshape(*input_shape, -1).contiguous()
        attn_output, _ = self.o_proj(attn_output)

        return attn_output


class Llama4VisionEncoderLayer(nn.Module):
    def __init__(
        self,
        config: Llama4VisionConfig,
340
        quant_config: QuantizationConfig | None,
341
        prefix: str = "",
342
        use_data_parallel: bool = False,
343
344
345
346
347
348
    ):
        super().__init__()
        self.hidden_size = config.hidden_size
        self.num_attention_heads = config.num_attention_heads
        self.intermediate_size = config.intermediate_size

349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
        self.self_attn = Llama4VisionAttention(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.self_attn",
            use_data_parallel=use_data_parallel,
        )
        self.mlp = Llama4VisionMLP(
            input_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            output_size=config.hidden_size,
            bias=True,
            output_activation=False,
            quant_config=quant_config,
            prefix=f"{prefix}.mlp",
            use_data_parallel=use_data_parallel,
        )
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384

        self.input_layernorm = nn.LayerNorm(config.hidden_size)
        self.post_attention_layernorm = nn.LayerNorm(config.hidden_size)

    def forward(
        self,
        hidden_state: torch.Tensor,
    ):
        # Self Attention
        residual = hidden_state
        hidden_state = self.input_layernorm(hidden_state)
        hidden_state = self.self_attn(hidden_state)
        hidden_state = residual + hidden_state

        # Feed forward
        residual = hidden_state
        hidden_state = self.post_attention_layernorm(hidden_state)
        hidden_state = self.mlp(hidden_state)
        hidden_state = residual + hidden_state

385
        outputs = (hidden_state,)
386
387
388
389
390
391
392
        return outputs


class Llama4VisionEncoder(nn.Module):
    def __init__(
        self,
        config: Llama4VisionConfig,
393
        quant_config: QuantizationConfig | None,
394
        prefix: str = "",
395
        use_data_parallel: bool = False,
396
397
398
    ):
        super().__init__()
        self.config = config
399
400
401
402
403
404
405
406
407
408
409
        self.layers = nn.ModuleList(
            [
                Llama4VisionEncoderLayer(
                    config,
                    quant_config=quant_config,
                    prefix=f"{prefix}.layers.{layer_idx}",
                    use_data_parallel=use_data_parallel,
                )
                for layer_idx in range(config.num_hidden_layers)
            ]
        )
410
411
412
413
414
415
416

    def forward(
        self,
        hidden_states: torch.Tensor,
    ) -> torch.Tensor:
        r"""
        Args:
417
            hidden_states: Input tensor of shape
418
                (batch_size, sequence_length, hidden_size).
419
                Hidden states from the model embeddings, representing
420
                the input tokens.
421
422
423
424
425
426
427
428
429
430
431
432
                associated vectors than the model's internal embedding
                lookup matrix.
        """

        for encoder_layer in self.layers:
            layer_outputs = encoder_layer(hidden_states)
            hidden_states = layer_outputs[0]

        return hidden_states


class Llama4UnfoldConvolution(nn.Module):
433
434
435
    def __init__(
        self,
        config: Llama4VisionConfig,
436
        quant_config: QuantizationConfig | None = None,
437
438
439
        prefix: str = "",
        use_data_parallel: bool = False,
    ):
440
441
442
443
        super().__init__()
        kernel_size = config.patch_size
        if isinstance(kernel_size, int):
            kernel_size = (kernel_size, kernel_size)
444
        self.unfold = torch.nn.Unfold(kernel_size=kernel_size, stride=config.patch_size)
445
446
447
448
449
450
451
452
453
        self.linear = ColumnParallelLinear(
            input_size=config.num_channels * kernel_size[0] * kernel_size[1],
            output_size=config.hidden_size,
            bias=False,
            gather_output=True,
            quant_config=quant_config,
            prefix=f"{prefix}.linear",
            disable_tp=use_data_parallel,
        )
454
455
456
457
458
459
460
461

    def forward(self, hidden_states: torch.Tensor) -> torch.Tensor:
        hidden_states = self.unfold(hidden_states)
        hidden_states = hidden_states.permute(0, 2, 1)
        hidden_states, _ = self.linear(hidden_states)
        return hidden_states


462
463
464
@support_torch_compile(
    dynamic_arg_dims={"images_flattened": 0}, enable_if=should_torch_compile_mm_vit
)
465
466
467
468
class Llama4VisionModel(nn.Module):
    def __init__(
        self,
        config: Llama4VisionConfig,
469
        quant_config: QuantizationConfig | None = None,
470
        prefix: str = "",
471
        use_data_parallel: bool = False,
472
473
474
475
476
477
478
479
    ):
        super().__init__()
        self.config = config
        self.image_size = config.image_size
        self.patch_size = config.patch_size
        self.hidden_size = config.hidden_size
        self.num_channels = config.num_channels

480
        self.num_patches = (self.image_size // self.patch_size) ** 2 + 1
481
482
483
484
485
        self.scale = config.hidden_size**-0.5

        self.patch_embedding = Llama4UnfoldConvolution(
            config,
            quant_config=quant_config,
486
487
488
            prefix=f"{prefix}.patch_embedding",
            use_data_parallel=use_data_parallel,
        )
489

490
        self.class_embedding = nn.Parameter(self.scale * torch.randn(self.hidden_size))
491
        self.positional_embedding_vlm = nn.Parameter(
492
493
            self.scale * torch.randn(self.num_patches, self.hidden_size)
        )
494
495
496
497
498
499

        # layer norms
        self.layernorm_pre = nn.LayerNorm(self.hidden_size, eps=1e-5)
        self.layernorm_post = nn.LayerNorm(self.hidden_size, eps=1e-5)

        # encoders
500
501
502
503
504
505
        self.model = Llama4VisionEncoder(
            config,
            quant_config=quant_config,
            prefix=f"{prefix}.model",
            use_data_parallel=use_data_parallel,
        )
506

507
        self.vision_adapter = Llama4VisionPixelShuffleMLP(
508
509
510
511
512
            config,
            quant_config,
            prefix=f"{prefix}.vision_adapter",
            use_data_parallel=use_data_parallel,
        )
513
514
515
516
517
518
519
520
521
522

    def forward(
        self,
        images_flattened: torch.Tensor,
    ) -> torch.Tensor:
        # Patch embedding
        hidden_state = self.patch_embedding(images_flattened)
        num_tiles, num_patches, hidden_dim = hidden_state.shape

        # Add cls token
523
524
525
        class_embedding = self.class_embedding.expand(
            hidden_state.shape[0], 1, hidden_state.shape[-1]
        )
526
527
528
529
530
531
532
533
534
535
536
        hidden_state = torch.cat([hidden_state, class_embedding], dim=1)
        num_patches += 1

        # Position embeddings
        hidden_state = hidden_state.reshape(
            num_tiles,
            1,
            num_patches,
            hidden_dim,
        )
        positional_embedding = self.positional_embedding_vlm.to(
537
538
            dtype=hidden_state.dtype, device=hidden_state.device
        )
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
        hidden_state = hidden_state + positional_embedding
        hidden_state = self.layernorm_pre(hidden_state)
        hidden_state = hidden_state.view(num_tiles, -1, hidden_dim)

        # Apply encoder
        hidden_state = self.model(hidden_state)
        hidden_state = self.layernorm_post(hidden_state)

        # Remove CLS token output
        hidden_state = hidden_state[:, :-1, :]

        # now, we use Llama4VisionPixelShuffle + mlp to project embeddings
        hidden_state = self.vision_adapter(hidden_state)

        return hidden_state


class Mllama4ProcessingInfo(BaseProcessingInfo):
    def __init__(self, ctx: InputProcessingContext) -> None:
        super().__init__(ctx)

    def get_hf_config(self) -> Llama4Config:
        return self.ctx.get_hf_config(Llama4Config)

    def get_hf_processor(self, **kwargs: object) -> Llama4Processor:
564
565
566
        return self.ctx.get_hf_processor(
            Llama4Processor, use_fast=kwargs.pop("use_fast", True), **kwargs
        )
567

568
    def get_supported_mm_limits(self) -> Mapping[str, int | None]:
569
570
571
        # Although vLLM can support more images from an infra capability
        # perspective, we do not recommend using >10 images in practice.
        return {"image": None}
572
573
574
575
576
577

    @staticmethod
    def get_patch_per_chunk(vision_config: Llama4VisionConfig) -> int:
        image_size = vision_config.image_size
        patch_size = vision_config.patch_size

578
579
580
        assert image_size % patch_size == 0, (
            f"chunk size {image_size} should be multiple of "
        )
581
582
583
        f"patch_size {patch_size}"

        ds_ratio = int(round(1.0 / (vision_config.pixel_shuffle_ratio**2)))
584
        return (image_size // patch_size) ** 2 // ds_ratio
585
586
587
588
589
590
591
592
593

    def get_max_num_tiles(self) -> int:
        image_processor = self.get_hf_processor().image_processor
        return image_processor.max_patches

    def get_image_size_with_most_features(self) -> ImageSize:
        vision_config = self.get_hf_config().vision_config
        image_size = vision_config.image_size
        # Result in the max possible feature size (h:w = 16:1)
594
        return ImageSize(height=self.get_max_num_tiles() * image_size, width=image_size)
595
596


597
class Mllama4MultiModalProcessor(BaseMultiModalProcessor[Mllama4ProcessingInfo]):
598
599
600
601
602
    def _call_hf_processor(
        self,
        prompt: str,
        mm_data: Mapping[str, object],
        mm_kwargs: Mapping[str, object],
603
        tok_kwargs: Mapping[str, object],
604
605
606
607
608
609
610
611
612
    ) -> BatchFeature:
        tokenizer = self.info.get_tokenizer()

        if mm_data is None:
            return tokenizer(prompt, add_special_tokens=False)  # exclude bos
        processed_outputs = super()._call_hf_processor(
            prompt=prompt,
            mm_data=mm_data,
            mm_kwargs=mm_kwargs,
613
            tok_kwargs=tok_kwargs,
614
615
616
617
618
619
620
        )

        processor = self.info.get_hf_processor(**mm_kwargs)
        image_processor = processor.image_processor
        vision_config = self.info.get_hf_config().vision_config

        if processed_outputs.get("pixel_values") is not None:
621
622
623
            assert "images" in mm_data, (
                "images expected to be in mm_data when pixel_values is present"
            )
624
625

            images = mm_data["images"]
626
627
628
629
630
            parsed_images = (
                self._get_data_parser()
                .parse_mm_data({"image": images})
                .get_items("image", ImageProcessorItems)
            )
631
632
633
634
635
636
637
638
639
640

            tile_size = vision_config.image_size
            possible_resolutions = find_supported_resolutions(
                max_num_chunks=self.info.get_max_num_tiles(),
                patch_size=SizeDict(height=tile_size, width=tile_size),
            )
            best_fit_sizes = [
                get_best_fit(
                    (image.size[1], image.size[0]),
                    torch.tensor(possible_resolutions),
641
                    resize_to_max_canvas=image_processor.resize_to_max_canvas,
642
643
                )
                for image in parsed_images
644
645
            ]
            # TODO tile height/width do not necessarily need to match
646
647
648
649
            aspect_ratios = [
                (image_size[0] // tile_size, image_size[1] // tile_size)
                for image_size in best_fit_sizes
            ]
650
            patches_per_image = [
651
                1 if r_h * r_w == 1 else 1 + r_h * r_w for (r_h, r_w) in aspect_ratios
652
653
            ]

654
            processed_outputs["aspect_ratios"] = torch.tensor(aspect_ratios)
655
            processed_outputs["patches_per_image"] = torch.tensor(patches_per_image)
656
657
658
659
660
661
662
663
664
665
666

        return processed_outputs

    def _get_mm_fields_config(
        self,
        hf_inputs: BatchFeature,
        hf_processor_mm_kwargs: Mapping[str, object],
    ) -> Mapping[str, MultiModalFieldConfig]:
        patches_per_image = hf_inputs.get("patches_per_image", torch.empty(0))
        return dict(
            pixel_values=MultiModalFieldConfig.flat_from_sizes(
667
668
                "image", patches_per_image
            ),
669
670
671
672
673
674
675
676
            patches_per_image=MultiModalFieldConfig.batched("image"),
            aspect_ratios=MultiModalFieldConfig.batched("image"),
        )

    def _get_prompt_updates(
        self,
        mm_items: MultiModalDataItems,
        hf_processor_mm_kwargs: Mapping[str, object],
677
        out_mm_kwargs: MultiModalKwargsItems,
678
    ) -> list[PromptUpdate]:
679
680
681
682
683
684
        config = self.info.get_hf_config()
        vision_config = config.vision_config

        num_patches_per_chunk = self.info.get_patch_per_chunk(vision_config)
        hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)
        image_token = hf_processor.image_token
685
        img_patch_token = hf_processor.img_patch_token
686
687

        def get_replacement(item_idx: int):
688
689
            out_item = out_mm_kwargs["image"][item_idx]
            aspect_ratio = out_item["aspect_ratios"].data
690
691

            repl = hf_processor._prompt_split_image(
692
                aspect_ratio=aspect_ratio,
693
694
695
696
                num_patches_per_chunk=num_patches_per_chunk,
            )

            return PromptUpdateDetails.select_text(repl, img_patch_token)
697
698
699
700
701
702
703
704
705
706
707

        return [
            PromptReplacement(
                modality="image",
                target=image_token,
                replacement=get_replacement,
            )
        ]


class Mllama4DummyInputsBuilder(BaseDummyInputsBuilder[Mllama4ProcessingInfo]):
708
709
710
711
712
713
714
715
716
    def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:
        num_images = mm_counts.get("image", 0)

        processor = self.info.get_hf_processor()
        image_token = processor.fake_image_token

        return image_token * num_images

    def get_dummy_mm_data(
717
718
719
        self,
        seq_len: int,
        mm_counts: Mapping[str, int],
720
        mm_options: Mapping[str, BaseDummyOptions] | None = None,
721
    ) -> MultiModalDataDict:
722
723
        num_images = mm_counts.get("image", 0)

724
        (target_width, target_height) = self.info.get_image_size_with_most_features()
725

726
727
        image_overrides = mm_options.get("image") if mm_options else None

728
        return {
729
730
731
732
733
734
            "image": self._get_dummy_images(
                width=target_width,
                height=target_height,
                num_images=num_images,
                overrides=image_overrides,
            )
735
736
737
738
739
740
741
742
        }


@MULTIMODAL_REGISTRY.register_processor(
    Mllama4MultiModalProcessor,
    info=Mllama4ProcessingInfo,
    dummy_inputs=Mllama4DummyInputsBuilder,
)
743
class Llama4ForConditionalGeneration(
744
745
746
747
748
749
    nn.Module,
    SupportsMultiModal,
    SupportsPP,
    MixtureOfExperts,
    SupportsEagle3,
    SupportsLoRA,
750
):
751
752
    packed_modules_mapping = {
        "qkv_proj": ["q_proj", "k_proj", "v_proj"],
753
        "gate_up_proj": ["gate_proj", "up_proj"],
754
755
    }

756
757
    supports_encoder_tp_data = True

758
    @classmethod
759
    def get_placeholder_str(cls, modality: str, i: int) -> str | None:
760
761
762
763
764
        if modality.startswith("image"):
            return "<|image|>"

        raise ValueError("Only image modality is supported")

765
766
767
768
769
    def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
        super().__init__()
        config = vllm_config.model_config.hf_config
        quant_config = vllm_config.quant_config
        multimodal_config = vllm_config.model_config.multimodal_config
770
771
        self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

772
        self.vllm_config = vllm_config
773
774
775
        self.config = config
        self.quant_config = quant_config
        self.multimodal_config = multimodal_config
776
        if multimodal_config.get_limit_per_prompt("image"):
777
778
779
780
781
782
783
784
785
786
787
788
789
            from vllm.compilation.backends import set_model_tag

            with (
                set_current_vllm_config(vllm_config),
                set_model_tag("Llama4VisionModel", is_encoder=True),
            ):
                self.vision_model = Llama4VisionModel(
                    config=config.vision_config,
                    quant_config=None,
                    prefix=maybe_prefix(prefix, "vision_model"),
                    use_data_parallel=self.use_data_parallel,
                )

790
            self.multi_modal_projector = Llama4MultiModalProjector(
791
792
793
                config=self.config,
                quant_config=None,
                prefix=maybe_prefix(prefix, "multi_modal_projector"),
794
            )
795
796
797
        else:
            self.vision_model = None
            self.multi_modal_projector = None
798
        self.language_model = initialize_model(
799
800
801
            vllm_config=vllm_config.with_hf_config(
                config.text_config, ["LlamaForCausalLM"]
            ),
802
803
804
805
806
            prefix=maybe_prefix(prefix, "language_model"),
            model_class=Llama4ForCausalLM,
        )

        self.make_empty_intermediate_tensors = (
807
808
            self.language_model.make_empty_intermediate_tensors
        )
809

810
811
812
813
814
815
816
817
818
819
820
        # Set MoE hyperparameters
        self.num_expert_groups = 1
        self.num_logical_experts = self.language_model.num_logical_experts
        self.num_physical_experts = self.language_model.num_physical_experts
        self.num_local_physical_experts = self.language_model.num_local_physical_experts
        self.num_routed_experts = self.language_model.num_routed_experts
        self.num_shared_experts = self.language_model.num_shared_experts
        self.num_redundant_experts = self.language_model.num_redundant_experts
        self.moe_layers = self.language_model.moe_layers
        self.num_moe_layers = len(self.moe_layers)

821
822
823
824
825
826
827
828
829
830
831
832
833
834
835
836
    def set_aux_hidden_state_layers(self, layers: tuple[int, ...]) -> None:
        """Set which layers should output auxiliary hidden states for EAGLE3."""
        # Delegate to underlying language model (Llama4ForCausalLM)
        assert hasattr(self.language_model, "set_aux_hidden_state_layers")
        self.language_model.set_aux_hidden_state_layers(layers)

    def get_eagle3_aux_hidden_state_layers(self) -> tuple[int, ...]:
        """Get the layer indices for auxiliary hidden state outputs.

        Note: The GPU model runner will override this with layers from
        the speculative config if available, providing dynamic configuration.
        """
        # Delegate to underlying language model (Llama4ForCausalLM)
        assert hasattr(self.language_model, "get_eagle3_aux_hidden_state_layers")
        return self.language_model.get_eagle3_aux_hidden_state_layers()

837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
    def set_eplb_state(
        self,
        expert_load_view: torch.Tensor,
        logical_to_physical_map: torch.Tensor,
        logical_replica_count: torch.Tensor,
    ):
        self.language_model.set_eplb_state(
            expert_load_view, logical_to_physical_map, logical_replica_count
        )
        self.expert_weights = self.language_model.expert_weights

    def update_physical_experts_metadata(
        self, num_physical_experts: int, num_local_physical_experts: int
    ):
        self.language_model.update_physical_experts_metadata(
            num_physical_experts, num_local_physical_experts
        )

855
    def _parse_and_validate_image_input(
856
        self, **kwargs: object
857
    ) -> Llama4ImagePatchInputs | None:
858
859
860
861
862
        # num_images, 1, num_chunks, channel, image_size, image_size
        pixel_values = kwargs.pop("pixel_values", None)
        if pixel_values is None:
            return None

863
        patches_per_image = kwargs.pop("patches_per_image")
864
        aspect_ratios = kwargs.pop("aspect_ratios")
865
866
867

        return Llama4ImagePatchInputs(
            type="pixel_values",
868
            pixel_values=pixel_values,
869
870
871
872
873
            patches_per_image=patches_per_image,
            aspect_ratios=aspect_ratios,
        )

    def _process_image_input(
874
875
        self, image_input: Llama4ImagePatchInputs
    ) -> MultiModalEmbeddings:
876
        assert self.vision_model and self.multi_modal_projector
877
        pixel_values = image_input["pixel_values"]
878
        patches_per_image = image_input["patches_per_image"].tolist()
879

880
881
882
        # shard image input
        if self.use_data_parallel:
            vision_embeddings_flat = run_dp_sharded_vision_model(
883
                pixel_values, self.vision_model
884
            )
885
        else:
886
            vision_embeddings_flat = self.vision_model(pixel_values)
887

888
        vision_embeddings_flat = self.multi_modal_projector(vision_embeddings_flat)
889
890
891
892
893

        return [
            img.flatten(0, 1)
            for img in vision_embeddings_flat.split(patches_per_image, dim=0)
        ]
894

895
896
897
    def get_language_model(self) -> torch.nn.Module:
        return self.language_model

898
    def embed_multimodal(self, **kwargs) -> MultiModalEmbeddings:
899
900
        image_input = self._parse_and_validate_image_input(**kwargs)
        if image_input is None:
901
            return []
902

903
904
905
906
        with (
            set_forward_context(None, self.vllm_config),
        ):
            return self._process_image_input(image_input)
907
908
909
910
911

    def forward(
        self,
        input_ids: torch.Tensor,
        positions: torch.Tensor,
912
913
        intermediate_tensors: IntermediateTensors | None = None,
        inputs_embeds: torch.Tensor | None = None,
914
        **kwargs: object,
915
    ) -> torch.Tensor | IntermediateTensors:
916
917
918
        if intermediate_tensors is not None:
            inputs_embeds = None

919
920
921
        return self.language_model(
            input_ids, positions, intermediate_tensors, inputs_embeds
        )
922
923
924
925

    def compute_logits(
        self,
        hidden_states: torch.Tensor,
926
    ) -> torch.Tensor | None:
927
        return self.language_model.compute_logits(hidden_states)
928
929
930

    def separate_weights(
        self,
931
        weights: Iterable[tuple[str, torch.Tensor]],
932
        prefix: str,
933
    ) -> tuple[Iterable[tuple[str, torch.Tensor]], Iterable[tuple[str, torch.Tensor]]]:
934
935
        weights1, weights2 = tee(weights, 2)

936
        def get_prefix_weights() -> Iterable[tuple[str, torch.Tensor]]:
937
938
939
940
            for name, data in weights1:
                if name.startswith(prefix):
                    yield (name, data)

941
        def get_other_weights() -> Iterable[tuple[str, torch.Tensor]]:
942
943
944
945
946
947
            for name, data in weights2:
                if not name.startswith(prefix):
                    yield (name, data)

        return get_prefix_weights(), get_other_weights()

948
949
950
951
952
953
954
955
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
    def _consolidate_qkv_weights(
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> Iterable[tuple[str, torch.Tensor]]:
        qkv_idx_mappings = {
            ".self_attn.q_proj": 0,
            ".self_attn.k_proj": 1,
            ".self_attn.v_proj": 2,
        }
        qkv_weights = {}
        for name, loaded_weight in weights:
            for weight_name, idx in qkv_idx_mappings.items():
                if weight_name not in name:
                    continue
                new_name = name.replace(weight_name, ".self_attn.qkv_proj")
                if new_name not in qkv_weights:
                    qkv_weights[new_name] = [None] * 3
                qkv_weights[new_name][idx] = loaded_weight
                break
            else:
                yield name, loaded_weight
        for key, weight in qkv_weights.items():
            qkv_weight = torch.cat(weight, dim=0)
            yield key, qkv_weight

972
973
974
    def _rename_weight_for_modelopt_checkpoint(self, name: str) -> str:
        """Rename weights from ModelOpt llama4 fp8 checkpoints to vLLM
        format."""
975
976
977
978
979
980
        if name.startswith("model.") or name.startswith("language_model.model."):
            renamed = (
                name.replace("model.", "language_model.model.", 1)
                if name.startswith("model.")
                else name
            )
981
            # Handle expert scale parameters with flat naming
982
983
984
            if "feed_forward.experts." in name and (
                "_input_scale" in name or "_weight_scale" in name
            ):
985
986
                # Map checkpoint naming to vLLM's expected naming
                if "down_proj_input_scale" in renamed:
987
                    return renamed.replace("down_proj_input_scale", "w2_input_scale")
988
                elif "down_proj_weight_scale" in renamed:
989
                    return renamed.replace("down_proj_weight_scale", "w2_weight_scale")
990
                elif "gate_up_proj_input_scale" in renamed:
991
992
993
                    return renamed.replace(
                        "gate_up_proj_input_scale", "w13_input_scale"
                    )
994
                elif "gate_up_proj_weight_scale" in renamed:
995
996
997
                    return renamed.replace(
                        "gate_up_proj_weight_scale", "w13_weight_scale"
                    )
998
999
1000
                return renamed

            # Handle attention scale parameters
1001
            elif "self_attn." in name and (".k_scale" in name or ".v_scale" in name):
1002
1003
1004
1005
1006
1007
1008
                if ".k_proj.k_scale" in renamed:
                    return renamed.replace(".k_proj.k_scale", ".attn.k_scale")
                elif ".v_proj.v_scale" in renamed:
                    return renamed.replace(".v_proj.v_scale", ".attn.v_scale")
                return renamed

            # Standard model.* to language_model.model.* renaming
1009
            return renamed
1010
1011

        elif name.startswith("lm_head.weight"):
1012
            return name.replace("lm_head.weight", "language_model.lm_head.weight")
1013
1014
1015
1016
1017
1018
1019
1020
1021
1022

        return name

    def _separate_and_rename_weights(
        self, weights: Iterable[tuple[str, torch.Tensor]]
    ) -> tuple[list[tuple[str, torch.Tensor]], list[tuple[str, torch.Tensor]]]:
        """Rename weights and separate them into language_model and other
        weights."""
        language_model_weights = []
        other_weights = []
1023

1024
1025
        for name, weight in weights:
            renamed = self._rename_weight_for_modelopt_checkpoint(name)
1026

1027
1028
1029
1030
1031
1032
1033
1034
            if renamed.startswith("language_model."):
                language_model_weights.append((renamed, weight))
            else:
                other_weights.append((renamed, weight))

        return language_model_weights, other_weights

    def _handle_expert_scale_broadcasting(
1035
        self, weights: list[tuple[str, torch.Tensor]], params_dict: dict
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047
    ) -> tuple[list[tuple[str, torch.Tensor]], set[str]]:
        """Handle expert scale parameters that need broadcasting.

        ModelOpt checkpoints use a single value tensor scalar for BMM style
        experts, vLLM expects the scale to be broadcasted across all experts.
        """
        regular_weights = []
        expert_scale_weights = []
        updated_params = set()

        for name, weight in weights:
            # Check if this is an expert scale parameter that needs broadcasting
1048
1049
1050
1051
1052
            if (
                "feed_forward.experts." in name
                and "scale" in name
                and ".shared_expert" not in name
            ):
1053
1054
                if name in params_dict:
                    param = params_dict[name]
1055
1056
1057
1058
1059
                    if (
                        hasattr(param, "data")
                        and param.data.numel() > 1
                        and weight.numel() == 1
                    ):
1060
1061
1062
1063
1064
1065
1066
1067
1068
1069
1070
                        # Broadcast single value to all experts
                        param.data.fill_(weight.item())
                        updated_params.add(name)
                        continue

                expert_scale_weights.append((name, weight))
            else:
                regular_weights.append((name, weight))

        return regular_weights, expert_scale_weights, updated_params

1071
1072
1073
1074
1075
1076
    def _load_other_weights(
        self,
        other_weights: Iterable[tuple[str, torch.Tensor]],
        params_dict: dict,
        stacked_params_mapping: list,
    ) -> set[str]:
1077
1078
        """Load non-language-model weights with stacking support."""
        updated_params = set()
1079

1080
1081
1082
        if self.use_data_parallel:
            other_weights = self._consolidate_qkv_weights(other_weights)

1083
        for name, loaded_weight in other_weights:
1084
            # Try stacked parameter mapping first
1085
            for param_name, weight_name, shard_id in stacked_params_mapping:
1086
                if weight_name not in name or self.use_data_parallel:
1087
1088
1089
1090
1091
1092
1093
1094
                    continue
                name = name.replace(weight_name, param_name)
                param = params_dict[name]
                updated_params.add(name)
                weight_loader = param.weight_loader
                weight_loader(param, loaded_weight, shard_id)
                break
            else:
1095
                # Use regular weight loading
1096
                param = params_dict[name]
1097
                weight_loader = getattr(param, "weight_loader", default_weight_loader)
1098
1099
                weight_loader(param, loaded_weight)
                updated_params.add(name)
1100
1101
1102

        return updated_params

1103
1104
1105
1106
    def get_expert_mapping(self) -> list[tuple[str, str, int, str]]:
        # Params for weights, fp8 weight scales, fp8 activation scales
        # (param_name, weight_name, expert_id, shard_id)
        return FusedMoE.make_expert_params_mapping(
1107
            self,
1108
1109
1110
1111
1112
1113
1114
            ckpt_gate_proj_name="gate_proj",
            ckpt_down_proj_name="down_proj",
            ckpt_up_proj_name="up_proj",
            num_experts=self.config.text_config.num_local_experts,
            num_redundant_experts=self.num_redundant_experts,
        )

1115
    def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
        stacked_params_mapping = [
            # (param_name, shard_name, shard_id)
            (".self_attn.qkv_proj", ".self_attn.q_proj", "q"),
            (".self_attn.qkv_proj", ".self_attn.k_proj", "k"),
            (".self_attn.qkv_proj", ".self_attn.v_proj", "v"),
            # Shared expert gate_up_proj stacking
            (".shared_expert.gate_up_proj", ".shared_expert.gate_proj", 0),
            (".shared_expert.gate_up_proj", ".shared_expert.up_proj", 1),
            # Feed forward gate_up_proj stacking (for non-MoE layers if any)
            (".feed_forward.gate_up_proj", ".feed_forward.gate_proj", 0),
            (".feed_forward.gate_up_proj", ".feed_forward.up_proj", 1),
        ]
        params_dict = dict(self.named_parameters())
        updated_params: set[str] = set()

        # Separate and rename weights
1132
1133
1134
        language_model_weights, other_weights = self._separate_and_rename_weights(
            weights
        )
1135

1136
1137
1138
1139
        # Skip loading vision model and projector if they're not initialized.
        if self.vision_model is None and self.multi_modal_projector is None:
            other_weights = []

1140
1141
        # Handle expert scale parameters
        regular_weights, expert_scale_weights, updated_params_from_experts = (
1142
1143
            self._handle_expert_scale_broadcasting(language_model_weights, params_dict)
        )
1144
1145
1146
1147
1148
1149
1150
1151
        updated_params.update(updated_params_from_experts)

        loader = AutoWeightsLoader(self)
        loaded_language_model_params = loader.load_weights(regular_weights)
        assert loaded_language_model_params is not None
        updated_params.update(loaded_language_model_params)

        if expert_scale_weights:
1152
            loaded_expert_scale_params = loader.load_weights(expert_scale_weights)
1153
1154
1155
1156
            if loaded_expert_scale_params:
                updated_params.update(loaded_expert_scale_params)

        updated_params.update(
1157
1158
            self._load_other_weights(other_weights, params_dict, stacked_params_mapping)
        )
1159

1160
        return updated_params
1161
1162
1163
1164
1165
1166
1167
1168
1169
1170

    def get_mm_mapping(self) -> MultiModelKeys:
        """
        Get the module prefix in multimodal models
        """
        return MultiModelKeys.from_string_field(
            language_model="language_model",
            connector="multi_modal_projector.",
            tower_model="vision_model.",
        )