@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / c2612371ad
Fork 0

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 13, 00:45 UTC

[Model] Add Gemma3 GGUF multimodal support (#27772)

Signed-off-by: Luciano Martins <lucianommartins@users.noreply.github.com> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Luciano Martins <lucianommartins@users.noreply.github.com> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Luciano Martins committed Nov 18, 2025 at 13:56 UTC c2612371ad76a966cbbc443da3f3f91a4f4a3138
14 files changed +751 -85
requirements/common.txt
+1 -1
@@ -30,7 +30,7 @@ filelock >= 3.16.1 # need to contain https://github.com/tox-dev/filelock/pull/31
30 partial-json-parser # used for parsing partial JSON outputs
31 pyzmq >= 25.0.0
32 msgspec
33 -gguf >= 0.13.0
33 +gguf >= 0.17.0
34 mistral_common[image] >= 1.8.5
35 opencv-python-headless >= 4.11.0 # required for video IO
36 pyyaml
tests/models/multimodal/generation/test_multimodal_gguf.py new
+115
@@ -0,0 +1,115 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +
4 +from typing import Literal, NamedTuple
5 +
6 +import pytest
7 +from huggingface_hub import hf_hub_download
8 +from pytest import MarkDecorator
9 +
10 +from tests.quantization.utils import is_quant_method_supported
11 +from vllm.assets.image import ImageAsset
12 +from vllm.utils.torch_utils import set_default_torch_num_threads
13 +
14 +from ....conftest import PromptImageInput, VllmRunner
15 +from ...utils import check_logprobs_close
16 +
17 +
18 +class GGUFMMTestConfig(NamedTuple):
19 + original_model: str
20 + gguf_repo: str
21 + gguf_backbone: str
22 + gguf_mmproj: str
23 + prompt: list[str]
24 + mm_data: dict[Literal["images"], PromptImageInput]
25 + max_model_len: int = 4096
26 + marks: list[MarkDecorator] = []
27 +
28 + @property
29 + def gguf_model(self):
30 + hf_hub_download(self.gguf_repo, filename=self.gguf_mmproj)
31 + return hf_hub_download(self.gguf_repo, filename=self.gguf_backbone)
32 +
33 +
34 +GEMMA3_CONFIG = GGUFMMTestConfig(
35 + original_model="google/gemma-3-4b-it",
36 + gguf_repo="google/gemma-3-4b-it-qat-q4_0-gguf",
37 + gguf_backbone="gemma-3-4b-it-q4_0.gguf",
38 + gguf_mmproj="mmproj-model-f16-4B.gguf",
39 + prompt=["<start_of_image>Describe this image in detail:"],
40 + mm_data={"images": [ImageAsset("stop_sign").pil_image]},
41 + marks=[pytest.mark.core_model],
42 +)
43 +
44 +MODELS_TO_TEST = [GEMMA3_CONFIG]
45 +
46 +
47 +def run_multimodal_gguf_test(
48 + vllm_runner: type[VllmRunner],
49 + model: GGUFMMTestConfig,
50 + dtype: str,
51 + max_tokens: int,
52 + num_logprobs: int,
53 +):
54 + # Run gguf model.
55 + with (
56 + set_default_torch_num_threads(1),
57 + vllm_runner(
58 + model_name=model.gguf_model,
59 + enforce_eager=True,
60 + tokenizer_name=model.original_model,
61 + dtype=dtype,
62 + max_model_len=model.max_model_len,
63 + ) as gguf_model,
64 + ):
65 + gguf_outputs = gguf_model.generate_greedy_logprobs(
66 + prompts=model.prompt,
67 + max_tokens=max_tokens,
68 + num_logprobs=num_logprobs,
69 + **model.mm_data,
70 + )
71 +
72 + # Run unquantized model.
73 + with vllm_runner(
74 + model_name=model.original_model,
75 + enforce_eager=True, # faster tests
76 + dtype=dtype,
77 + max_model_len=model.max_model_len,
78 + ) as original_model:
79 + original_outputs = original_model.generate_greedy_logprobs(
80 + prompts=model.prompt,
81 + max_tokens=max_tokens,
82 + num_logprobs=num_logprobs,
83 + **model.mm_data,
84 + )
85 +
86 + check_logprobs_close(
87 + outputs_0_lst=original_outputs,
88 + outputs_1_lst=gguf_outputs,
89 + name_0="original",
90 + name_1="gguf",
91 + )
92 +
93 +
94 +@pytest.mark.skipif(
95 + not is_quant_method_supported("gguf"),
96 + reason="gguf is not supported on this GPU type.",
97 +)
98 +@pytest.mark.parametrize(
99 + "model",
100 + [
101 + pytest.param(test_config, marks=test_config.marks)
102 + for test_config in MODELS_TO_TEST
103 + ],
104 +)
105 +@pytest.mark.parametrize("dtype", ["bfloat16"])
106 +@pytest.mark.parametrize("max_tokens", [32])
107 +@pytest.mark.parametrize("num_logprobs", [10])
108 +def test_models(
109 + vllm_runner: type[VllmRunner],
110 + model: GGUFMMTestConfig,
111 + dtype: str,
112 + max_tokens: int,
113 + num_logprobs: int,
114 +) -> None:
115 + run_multimodal_gguf_test(vllm_runner, model, dtype, max_tokens, num_logprobs)
tests/models/quantization/test_gguf.py
+8 -1
@@ -78,6 +78,12 @@ DOLPHIN_CONFIG = GGUFTestConfig(
78 gguf_filename="tinydolphin-2.8-1.1b.Q6_K.gguf",
79 )
80
81 +GEMMA3_CONFIG = GGUFTestConfig(
82 + original_model="google/gemma-3-270m-it",
83 + gguf_repo="ggml-org/gemma-3-270m-it-qat-GGUF",
84 + gguf_filename="gemma-3-270m-it-qat-Q4_0.gguf",
85 +)
86 +
87 MODELS = [
88 # LLAMA_CONFIG, # broken: https://github.com/vllm-project/vllm/issues/19458
89 QWEN2_CONFIG,
@@ -85,6 +91,7 @@ MODELS = [
91 GPT2_CONFIG,
92 STABLELM_CONFIG,
93 DOLPHIN_CONFIG,
94 + GEMMA3_CONFIG,
95 # STARCODER_CONFIG, # broken
96 ]
97
@@ -148,7 +155,7 @@ def check_model_outputs(
155 "model",
156 [pytest.param(test_config, marks=test_config.marks) for test_config in MODELS],
157 )
151 -@pytest.mark.parametrize("dtype", ["half"])
158 +@pytest.mark.parametrize("dtype", ["bfloat16"])
159 @pytest.mark.parametrize("max_tokens", [32])
160 @pytest.mark.parametrize("num_logprobs", [5])
161 @pytest.mark.parametrize("tp_size", [1])
vllm/config/model.py
+19 -1
@@ -33,10 +33,14 @@ from vllm.transformers_utils.config import (
33 try_get_generation_config,
34 try_get_safetensors_metadata,
35 try_get_tokenizer_config,
36 + uses_custom_attention_masks,
37 uses_mrope,
38 )
39 +from vllm.transformers_utils.gguf_utils import (
40 + maybe_patch_hf_config_from_gguf,
41 +)
42 from vllm.transformers_utils.runai_utils import ObjectStorageModel, is_runai_obj_uri
39 -from vllm.transformers_utils.utils import maybe_model_redirect
43 +from vllm.transformers_utils.utils import check_gguf_file, maybe_model_redirect
44 from vllm.utils.import_utils import LazyLoader
45 from vllm.utils.torch_utils import common_broadcastable_dtype
46
@@ -450,6 +454,12 @@ class ModelConfig:
454 self.model = maybe_model_redirect(self.model)
455 # The tokenizer is consistent with the model by default.
456 if self.tokenizer is None:
457 + if check_gguf_file(self.model):
458 + raise ValueError(
459 + "Using a tokenizer is mandatory when loading a GGUF model. "
460 + "Please specify the tokenizer path or name using the "
461 + "--tokenizer argument."
462 + )
463 self.tokenizer = self.model
464 if self.tokenizer_revision is None:
465 self.tokenizer_revision = self.revision
@@ -508,6 +518,10 @@ class ModelConfig:
518 hf_overrides_kw=hf_overrides_kw,
519 hf_overrides_fn=hf_overrides_fn,
520 )
521 + hf_config = maybe_patch_hf_config_from_gguf(
522 + self.model,
523 + hf_config,
524 + )
525
526 self.hf_config = hf_config
527 if dict_overrides:
@@ -1605,6 +1619,10 @@ class ModelConfig:
1619 def uses_mrope(self) -> bool:
1620 return uses_mrope(self.hf_config)
1621
1622 + @property
1623 + def uses_custom_attention_masks(self) -> bool:
1624 + return uses_custom_attention_masks(self.hf_config)
1625 +
1626 @property
1627 def is_multimodal_model(self) -> bool:
1628 return self.multimodal_config is not None
vllm/model_executor/layers/quantization/gguf.py
+62 -5
@@ -1,7 +1,8 @@
1 # SPDX-License-Identifier: Apache-2.0
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4 -from collections.abc import Callable
4 +from collections.abc import Callable, Mapping
5 +from types import MappingProxyType
6 from typing import Any, Optional
7
8 import gguf
@@ -26,7 +27,11 @@ from vllm.model_executor.layers.quantization.base_config import (
27 QuantizationConfig,
28 QuantizeMethodBase,
29 )
29 -from vllm.model_executor.layers.vocab_parallel_embedding import VocabParallelEmbedding
30 +from vllm.model_executor.layers.vocab_parallel_embedding import (
31 + UnquantizedEmbeddingMethod,
32 + VocabParallelEmbedding,
33 +)
34 +from vllm.model_executor.models.utils import WeightsMapper
35 from vllm.model_executor.utils import set_weight_attrs
36 from vllm.utils.torch_utils import direct_register_custom_op
37
@@ -65,18 +70,70 @@ class GGUFConfig(QuantizationConfig):
70 self, layer: torch.nn.Module, prefix: str
71 ) -> Optional["QuantizeMethodBase"]:
72 if isinstance(layer, LinearBase):
68 - if is_layer_skipped_gguf(prefix, self.unquantized_modules):
73 + if is_layer_skipped_gguf(
74 + prefix, self.unquantized_modules, self.packed_modules_mapping
75 + ):
76 return UnquantizedLinearMethod()
77 return GGUFLinearMethod(self)
78 elif isinstance(layer, VocabParallelEmbedding):
79 + if is_layer_skipped_gguf(
80 + prefix, self.unquantized_modules, self.packed_modules_mapping
81 + ):
82 + return UnquantizedEmbeddingMethod()
83 return GGUFEmbeddingMethod(self)
84 elif isinstance(layer, FusedMoE):
85 return GGUFMoEMethod(self, layer.moe_config)
86 return None
87
88 + def apply_vllm_mapper(self, hf_to_vllm_mapper: "WeightsMapper"):
89 + """
90 + Interface for models to update module names referenced in
91 + quantization configs in order to reflect the vllm model structure
92 +
93 + :param hf_to_vllm_mapper: maps from hf model structure (the assumed
94 + structure of the qconfig) to vllm model structure
95 + """
96 + if self.unquantized_modules is not None:
97 + self.unquantized_modules = hf_to_vllm_mapper.apply_list(
98 + self.unquantized_modules
99 + )
100 +
101 +
102 +def is_layer_skipped_gguf(
103 + prefix: str,
104 + unquantized_modules: list[str],
105 + fused_mapping: Mapping[str, list[str]] = MappingProxyType({}),
106 +):
107 + # Fused layers like gate_up_proj or qkv_proj will not be fused
108 + # in the safetensors checkpoint. So, we convert the name
109 + # from the fused version to unfused + check to make sure that
110 + # each shard of the fused layer has the same scheme.
111 + proj_name = prefix.split(".")[-1]
112 + if proj_name in fused_mapping:
113 + shard_prefixes = [
114 + prefix.replace(proj_name, shard_proj_name)
115 + for shard_proj_name in fused_mapping[proj_name]
116 + ]
117 +
118 + is_skipped = None
119 + for shard_prefix in shard_prefixes:
120 + is_shard_skipped = any(
121 + shard_prefix in module_name for module_name in unquantized_modules
122 + )
123 +
124 + if is_skipped is None:
125 + is_skipped = is_shard_skipped
126 + elif is_shard_skipped != is_skipped:
127 + raise ValueError(
128 + f"Detected some but not all shards of {prefix} "
129 + "are quantized. All shards of fused layers "
130 + "to have the same precision."
131 + )
132 + else:
133 + is_skipped = any(module_name in prefix for module_name in unquantized_modules)
134
78 -def is_layer_skipped_gguf(prefix: str, unquantized_modules: list[str]):
79 - return any(module_name in prefix for module_name in unquantized_modules)
135 + assert is_skipped is not None
136 + return is_skipped
137
138
139 UNQUANTIZED_TYPES = {WeightType.F32, WeightType.F16, WeightType.BF16}
vllm/model_executor/model_loader/gguf_loader.py
+173 -13
@@ -7,10 +7,11 @@ import gguf
7 import torch
8 import torch.nn as nn
9 from huggingface_hub import hf_hub_download
10 -from transformers import AutoModelForCausalLM
10 +from transformers import AutoModelForCausalLM, AutoModelForImageTextToText
11
12 from vllm.config import ModelConfig, VllmConfig
13 from vllm.config.load import LoadConfig
14 +from vllm.logger import init_logger
15 from vllm.model_executor.model_loader.base_loader import BaseModelLoader
16 from vllm.model_executor.model_loader.utils import (
17 initialize_model,
@@ -21,8 +22,11 @@ from vllm.model_executor.model_loader.weight_utils import (
22 get_gguf_weight_type_map,
23 gguf_quant_weights_iterator,
24 )
25 +from vllm.transformers_utils.gguf_utils import detect_gguf_multimodal
26 from vllm.utils.torch_utils import set_default_torch_dtype
27
28 +logger = init_logger(__name__)
29 +
30
31 class GGUFModelLoader(BaseModelLoader):
32 """
@@ -67,7 +71,15 @@ class GGUFModelLoader(BaseModelLoader):
71 https://github.com/ggerganov/ggml/blob/master/docs/gguf.md for details.
72 """
73 config = model_config.hf_config
74 + # Get text config to handle both nested (multimodal) and flat
75 + # (text-only) config structures. For multimodal models like
76 + # Gemma3Config, this returns config.text_config. For text-only
77 + # models, this returns config itself.
78 + text_config = config.get_text_config()
79 model_type = config.model_type
80 + is_multimodal = (
81 + hasattr(config, "vision_config") and config.vision_config is not None
82 + )
83 gguf_to_hf_name_map = {}
84 # hack: ggufs have a different name than transformers
85 if model_type == "cohere":
@@ -115,24 +127,167 @@ class GGUFModelLoader(BaseModelLoader):
127 break
128 if arch is None:
129 raise RuntimeError(f"Unknown gguf model_type: {model_type}")
118 - num_layers = config.num_hidden_layers
119 - name_map = gguf.get_tensor_name_map(arch, num_layers)
130 + text_num_layers = text_config.num_hidden_layers
131 + text_name_map = gguf.get_tensor_name_map(arch, text_num_layers)
132 +
133 + if is_multimodal:
134 + mm_proj_arch = gguf.MODEL_ARCH.MMPROJ
135 + vision_num_layers = config.vision_config.num_hidden_layers
136 + vision_name_map = gguf.get_tensor_name_map(mm_proj_arch, vision_num_layers)
137 + else:
138 + vision_name_map = None
139 +
140 + # Create dummy model to extract parameter names
141 + # For multimodal: use AutoModelForImageTextToText to get
142 + # language + vision + projector params
143 + # For text-only: use AutoModelForCausalLM to get language model params
144 + auto_cls = (
145 + AutoModelForImageTextToText if is_multimodal else AutoModelForCausalLM
146 + )
147 with torch.device("meta"):
121 - dummy_model = AutoModelForCausalLM.from_config(
148 + dummy_model = auto_cls.from_config(
149 config, trust_remote_code=model_config.trust_remote_code
150 )
151 +
152 state_dict = dummy_model.state_dict()
153 + if hf_checkpoint_map := getattr(
154 + dummy_model, "_checkpoint_conversion_mapping", None
155 + ):
156 +
157 + def revert_hf_rename(name: str) -> str:
158 + for original_name, hf_name in hf_checkpoint_map.items():
159 + if hf_name in name:
160 + name = name.replace(hf_name, original_name).lstrip("^")
161 + return name
162 +
163 + state_dict = {
164 + revert_hf_rename(name): tensor for name, tensor in state_dict.items()
165 + }
166 +
167 + def find_hf_name_in_tensor_map(hf_name: str) -> str | None:
168 + """
169 + Map HuggingFace parameter name to GGUF tensor name.
170 +
171 + This function handles the mismatch between HF parameter naming
172 + conventions and gguf-py's expected format:
173 + 1. Strips 'model.' prefix (common in multimodal models)
174 + 2. Converts '_weight' suffix to '.weight' (Gemma3 compatibility)
175 + 3. Searches vision_name_map for multimodal parameters
176 + 4. Falls back to text_name_map for language model parameters
177 +
178 + Args:
179 + hf_name: Full HuggingFace parameter name (e.g.,
180 + 'model.multi_modal_projector.mm_soft_emb_norm.weight')
181 +
182 + Returns:
183 + GGUF tensor name with suffix (e.g., 'mm.soft_emb_norm.weight')
184 + or None if no mapping found
185 + """
186 + # Strip 'language_model.' prefix for multimodal models - gguf-py
187 + # tensor mappings expect parameter names without this prefix.
188 + # Note: 'model.' prefix should be KEPT for text-only models as
189 + # gguf-py expects it.
190 + if hf_name.startswith("language_model."):
191 + hf_name = hf_name[15:] # Remove 'language_model.'
192 +
193 + # Parse parameter name and suffix
194 + if hf_name.endswith((".weight", ".bias")):
195 + base_name, suffix = hf_name.rsplit(".", 1)
196 + else:
197 + base_name, suffix = hf_name, ""
198 + # Handle '_weight' suffix (Gemma3 naming: parameter ends with
199 + # '_weight' instead of '.weight')
200 + if base_name.endswith("_weight"):
201 + base_name = base_name[:-7] # Remove '_weight'
202 + suffix = "weight"
203 +
204 + gguf_name = None
205 + # Priority 1: Search vision/projector parameters for multimodal models
206 + if vision_name_map is not None:
207 + gguf_name = vision_name_map.get_name(base_name)
208 +
209 + # Priority 2: Search text backbone parameters
210 + if gguf_name is None:
211 + gguf_name = text_name_map.get_name(base_name)
212 +
213 + if gguf_name is None:
214 + return None
215
216 + return gguf_name + "." + suffix
217 +
218 + # Build mapping and track unmapped parameters
219 + unmapped_params = []
220 for hf_name in state_dict:
127 - name, suffix = hf_name.rsplit(".", 1)
128 - gguf_name = name_map.get_name(name)
129 - gguf_to_hf_name_map[f"{gguf_name}.{suffix}"] = hf_name
221 + gguf_name_with_suffix = find_hf_name_in_tensor_map(hf_name)
222 +
223 + # Track mapping success
224 + if gguf_name_with_suffix is not None:
225 + gguf_to_hf_name_map[gguf_name_with_suffix] = hf_name
226 + logger.debug("Mapped GGUF %s → HF %s", gguf_name_with_suffix, hf_name)
227 + elif hf_name not in gguf_to_hf_name_map.values():
228 + # Parameter not in manual overrides either
229 + unmapped_params.append(hf_name)
230 +
231 + # All parameters must be mapped: both vision/projector and backbone
232 + if unmapped_params:
233 + raise RuntimeError(
234 + f"Failed to map GGUF parameters "
235 + f"({len(unmapped_params)}): "
236 + f"{unmapped_params}"
237 + )
238 return gguf_to_hf_name_map
239
240 + def _get_gguf_weight_type(
241 + self,
242 + model_config: ModelConfig,
243 + model_name_or_path: str,
244 + gguf_to_hf_name_map: dict[str, str],
245 + ) -> dict[str, str]:
246 + weight_type_map = get_gguf_weight_type_map(
247 + model_config.model, gguf_to_hf_name_map
248 + )
249 + is_multimodal = hasattr(model_config.hf_config, "vision_config")
250 + if is_multimodal:
251 + mmproj_file = detect_gguf_multimodal(model_name_or_path)
252 + assert mmproj_file is not None, (
253 + "Could not find mm_proj file for multimodal GGUF model"
254 + )
255 + logger.info("Loading extra mm_proj weights from %s...", mmproj_file)
256 + mm_proj_weight_type_map = get_gguf_weight_type_map(
257 + mmproj_file, gguf_to_hf_name_map
258 + )
259 + weight_type_map.update(mm_proj_weight_type_map)
260 + return weight_type_map
261 +
262 def _get_weights_iterator(
133 - self, model_name_or_path: str, gguf_to_hf_name_map: dict[str, str]
263 + self,
264 + model_config: ModelConfig,
265 + model_name_or_path: str,
266 + gguf_to_hf_name_map: dict[str, str],
267 ) -> Generator[tuple[str, torch.Tensor], None, None]:
135 - return gguf_quant_weights_iterator(model_name_or_path, gguf_to_hf_name_map)
268 + """
269 + Iterate over GGUF model weights, loading from both main model file and
270 + mmproj.gguf for multimodal Gemma3 models.
271 +
272 + For Gemma3 multimodal GGUF models:
273 + - Main file (gemma-3-*.gguf): Language model weights (model.*)
274 + - mmproj file (mmproj*.gguf): Vision tower + projector weights (v.*, mm.*)
275 +
276 + Yields:
277 + Tuples of (parameter_name, tensor) for all model weights
278 + """
279 + hf_config = model_config.hf_config
280 + is_multimodal = hasattr(hf_config, "vision_config")
281 +
282 + if is_multimodal:
283 + # Load mm_proj (mm_encoder + projector) for multimodal weights
284 + mmproj_file = detect_gguf_multimodal(model_name_or_path)
285 + assert mmproj_file is not None, (
286 + "Could not find mm_proj file for multimodal GGUF model"
287 + )
288 + yield from gguf_quant_weights_iterator(mmproj_file, gguf_to_hf_name_map)
289 +
290 + yield from gguf_quant_weights_iterator(model_name_or_path, gguf_to_hf_name_map)
291
292 def download_model(self, model_config: ModelConfig) -> None:
293 self._prepare_weights(model_config.model)
@@ -141,7 +296,7 @@ class GGUFModelLoader(BaseModelLoader):
296 local_model_path = self._prepare_weights(model_config.model)
297 gguf_weights_map = self._get_gguf_weights_map(model_config)
298 model.load_weights(
144 - self._get_weights_iterator(local_model_path, gguf_weights_map)
299 + self._get_weights_iterator(model_config, local_model_path, gguf_weights_map)
300 )
301
302 def load_model(
@@ -156,14 +311,19 @@ class GGUFModelLoader(BaseModelLoader):
311 ):
312 model_config.hf_config.update({"tie_word_embeddings": True})
313
159 - weight_type_map = get_gguf_weight_type_map(model_config.model, gguf_weights_map)
160 -
314 + weight_type_map = self._get_gguf_weight_type(
315 + model_config, local_model_path, gguf_weights_map
316 + )
317 # filter out unquantized modules to skip
318 unquant_names = [
319 name.removesuffix(".weight")
320 for name, weight_type in weight_type_map.items()
165 - if weight_type == "F32" and name.endswith(".weight")
321 + if weight_type in ("F32", "F16", "BF16") and name.endswith(".weight")
322 ]
323 + logger.debug(
324 + "GGUF unquantized modules: %s",
325 + unquant_names,
326 + )
327 vllm_config.quant_config.unquantized_modules.extend(unquant_names)
328
329 target_device = torch.device(device_config.device)
vllm/model_executor/model_loader/weight_utils.py
+7 -3
@@ -836,7 +836,11 @@ def gguf_quant_weights_iterator(
836 ) -> Generator[tuple[str, torch.Tensor], None, None]:
837 """
838 Iterate over the quant weights in the model gguf files and convert
839 - them to torch tensors
839 + them to torch tensors.
840 + Be careful of the order of yielding weight types and weights data,
841 + we have to yield all weight types first before yielding any weights.
842 + Otherwise it would cause issue when loading weights with for packed
843 + layer with different quant types.
844 """
845
846 reader = gguf.GGUFReader(gguf_file)
@@ -846,7 +850,7 @@ def gguf_quant_weights_iterator(
850 weight_type = tensor.tensor_type
851 name = gguf_to_hf_name_map[tensor.name]
852
849 - if weight_type.name != "F32":
853 + if weight_type.name not in ("F32", "BF16", "F16"):
854 weight_type_name = name.replace("weight", "qweight_type")
855 weight_type = torch.tensor(weight_type)
856 yield weight_type_name, weight_type
@@ -856,7 +860,7 @@ def gguf_quant_weights_iterator(
860 weight = tensor.data
861 weight_type = tensor.tensor_type
862 name = gguf_to_hf_name_map[tensor.name]
859 - if weight_type.name != "F32":
863 + if weight_type.name not in ("F32", "BF16", "F16"):
864 name = name.replace("weight", "qweight")
865 param = torch.tensor(weight)
866 yield name, param
vllm/model_executor/models/gemma3_mm.py
+116 -56
@@ -2,7 +2,7 @@
2 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 import math
4 from collections.abc import Iterable, Mapping, Sequence
5 -from typing import Annotated, Any, Literal, TypeAlias
5 +from typing import Annotated, Any, Literal
6
7 import torch
8 from torch import nn
@@ -20,12 +20,7 @@ from vllm.multimodal.inputs import (
20 MultiModalFieldConfig,
21 MultiModalKwargsItems,
22 )
23 -from vllm.multimodal.parse import (
24 - ImageEmbeddingItems,
25 - ImageProcessorItems,
26 - ImageSize,
27 - MultiModalDataItems,
28 -)
23 +from vllm.multimodal.parse import ImageProcessorItems, ImageSize, MultiModalDataItems
24 from vllm.multimodal.processing import (
25 BaseMultiModalProcessor,
26 BaseProcessingInfo,
@@ -76,15 +71,7 @@ class Gemma3ImagePixelInputs(TensorSchema):
71 num_patches: Annotated[torch.Tensor, TensorShape("bn")]
72
73
79 -class Gemma3ImageEmbeddingInputs(TensorSchema):
80 - type: Literal["image_embeds"] = "image_embeds"
81 - image_embeds: Annotated[
82 - torch.Tensor,
83 - TensorShape("ni", "nf", "hs"),
84 - ]
85 -
86 -
87 -Gemma3ImageInputs: TypeAlias = Gemma3ImagePixelInputs | Gemma3ImageEmbeddingInputs
74 +Gemma3ImageInputs = Gemma3ImagePixelInputs
75
76
77 class Gemma3ProcessingInfo(BaseProcessingInfo):
@@ -191,9 +178,8 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):
178 def get_image_repl(
179 self,
180 *,
194 - image_width: int | None,
195 - image_height: int | None,
196 - num_crops: int | None = None,
181 + image_width: int,
182 + image_height: int,
183 processor: Gemma3Processor | None,
184 ) -> PromptUpdateDetails[str]:
185 if processor is None:
@@ -201,13 +187,11 @@ class Gemma3ProcessingInfo(BaseProcessingInfo):
187
188 boi_token = processor.boi_token
189
204 - if num_crops is None:
205 - assert image_width is not None and image_height is not None
206 - num_crops = self.get_num_crops(
207 - image_width=image_width,
208 - image_height=image_height,
209 - processor=processor,
210 - )
190 + num_crops = self.get_num_crops(
191 + image_width=image_width,
192 + image_height=image_height,
193 + processor=processor,
194 + )
195
196 if num_crops == 0:
197 image_text = boi_token
@@ -337,7 +321,6 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
321 return dict(
322 pixel_values=MultiModalFieldConfig.flat_from_sizes("image", num_patches),
323 num_patches=MultiModalFieldConfig.batched("image"),
340 - image_embeds=MultiModalFieldConfig.batched("image"),
324 )
325
326 def _get_prompt_updates(
@@ -350,19 +333,7 @@ class Gemma3MultiModalProcessor(BaseMultiModalProcessor[Gemma3ProcessingInfo]):
333 image_token = hf_processor.boi_token
334
335 def get_replacement_gemma3(item_idx: int):
353 - images = mm_items.get_items(
354 - "image", (ImageEmbeddingItems, ImageProcessorItems)
355 - )
356 -
357 - if isinstance(images, ImageEmbeddingItems):
358 - # For image embedding inputs, only support no crops cases
359 - # since it's not supported in hf processor anyway
360 - return self.info.get_image_repl(
361 - image_width=None,
362 - image_height=None,
363 - num_crops=0,
364 - processor=hf_processor,
365 - )
336 + images = mm_items.get_items("image", ImageProcessorItems)
337
338 image_size = images.get_image_size(item_idx)
339 return self.info.get_image_repl(
@@ -586,19 +557,17 @@ class Gemma3ForConditionalGeneration(
557 pixel_values = kwargs.pop("pixel_values", None)
558 num_patches = kwargs.pop("num_patches", None)
559 image_embeds = kwargs.pop("image_embeds", None)
560 + assert image_embeds is None, "Gemma3 does not support image_embeds."
561 + if pixel_values is None:
562 + return None
563
590 - if pixel_values is not None:
591 - image_size = self.config.vision_config.image_size
592 - return Gemma3ImagePixelInputs(
593 - pixel_values=pixel_values,
594 - num_patches=num_patches,
595 - resolve_bindings={"h": image_size, "w": image_size},
596 - )
597 - elif image_embeds is not None:
598 - return Gemma3ImageEmbeddingInputs(
599 - image_embeds=image_embeds,
600 - type="image_embeds",
601 - )
564 + image_size = self.config.vision_config.image_size
565 +
566 + return Gemma3ImagePixelInputs(
567 + pixel_values=pixel_values,
568 + num_patches=num_patches,
569 + resolve_bindings={"h": image_size, "w": image_size},
570 + )
571
572 def _image_pixels_to_features(
573 self,
@@ -610,9 +579,7 @@ class Gemma3ForConditionalGeneration(
579 def _process_image_input(
580 self,
581 image_input: Gemma3ImageInputs,
613 - ) -> torch.Tensor | list[torch.Tensor]:
614 - if image_input["type"] == "image_embeds":
615 - return image_input["image_embeds"]
582 + ) -> list[torch.Tensor]:
583 assert self.vision_tower is not None
584
585 pixel_values = image_input["pixel_values"]
@@ -629,13 +596,33 @@ class Gemma3ForConditionalGeneration(
596 def get_language_model(self) -> torch.nn.Module:
597 return self.language_model
598
632 - def embed_multimodal(self, **kwargs: object) -> MultiModalEmbeddings:
599 + def get_multimodal_embeddings(self, **kwargs: object) -> MultiModalEmbeddings:
600 image_input = self._parse_and_validate_image_input(**kwargs)
601 if image_input is None:
602 return []
603
604 return self._process_image_input(image_input)
605
606 + def embed_input_ids(
607 + self,
608 + input_ids: torch.Tensor,
609 + multimodal_embeddings: MultiModalEmbeddings | None = None,
610 + *,
611 + is_multimodal: torch.Tensor | None = None,
612 + handle_oov_mm_token: bool = True,
613 + ) -> torch.Tensor:
614 + # Early return for text-only inference (no multimodal data)
615 + if multimodal_embeddings is None or is_multimodal is None:
616 + return super().embed_input_ids(input_ids)
617 +
618 + # Use interface default with OOV handling enabled
619 + return super().embed_input_ids(
620 + input_ids,
621 + multimodal_embeddings=multimodal_embeddings,
622 + is_multimodal=is_multimodal,
623 + handle_oov_mm_token=handle_oov_mm_token,
624 + )
625 +
626 def forward(
627 self,
628 input_ids: torch.Tensor,
@@ -657,6 +644,79 @@ class Gemma3ForConditionalGeneration(
644
645 return hidden_states
646
647 + def generate_attention_masks(
648 + self,
649 + input_ids: torch.Tensor,
650 + positions: torch.Tensor,
651 + mask_dtype: torch.dtype,
652 + ) -> dict[str, Any]:
653 + """Generate custom attention masks for Gemma3 multimodal inputs.
654 +
655 + This is called by V1 engine's gpu_model_runner during preprocessing
656 + to generate attention masks that allow bidirectional attention between
657 + image tokens while maintaining causal attention for text.
658 + """
659 + # NOTE(woosuk): Here, we distinguish the sequences by the position id 0.
660 + # This is a HACK. Fix this.
661 + start_indices = (positions == 0).cpu().nonzero()
662 + num_seqs = len(start_indices)
663 + seq_lens = []
664 + for i in range(num_seqs):
665 + start_idx = start_indices[i]
666 + end_idx = start_indices[i + 1] if i < num_seqs - 1 else len(input_ids)
667 + seq_lens.append(end_idx - start_idx)
668 +
669 + global_attn_masks = []
670 + local_attn_masks = []
671 + start_idx = 0
672 + for seq_idx, seq_len in enumerate(seq_lens):
673 + end_idx = start_idx + seq_len
674 + input_token_ids = input_ids[start_idx:end_idx]
675 +
676 + # Find image token positions
677 + img_pos = input_token_ids == self.config.image_token_index
678 +
679 + start_idx = end_idx
680 +
681 + # Create a global causal mask
682 + global_attn_mask = torch.empty(
683 + 1,
684 + 1,
685 + seq_len,
686 + seq_len,
687 + dtype=mask_dtype,
688 + device=input_ids.device,
689 + )
690 + global_attn_mask.fill_(float("-inf"))
691 + # Fill the lower triangle with 0 (causal attention)
692 + global_attn_mask = global_attn_mask.triu(diagonal=1)
693 +
694 + # Enable bidirectional attention between image tokens
695 + img_mask = torch.zeros_like(global_attn_mask)
696 + img_mask[:, :, :, img_pos] += 1
697 + img_mask[:, :, img_pos, :] += 1
698 + global_attn_mask = torch.where(img_mask == 2, 0, global_attn_mask)
699 + global_attn_masks.append(global_attn_mask)
700 +
701 + # GGUF compatibility: config might be Gemma3TextConfig directly
702 + text_config = getattr(self.config, "text_config", self.config)
703 + sliding_window = text_config.sliding_window
704 + if sliding_window is not None:
705 + # Create a local causal mask with sliding window (1024)
706 + local_attn_mask = torch.ones_like(global_attn_mask)
707 + local_attn_mask = torch.tril(local_attn_mask, diagonal=-sliding_window)
708 + local_attn_mask = torch.where(
709 + local_attn_mask == 0, global_attn_mask, float("-inf")
710 + )
711 + local_attn_masks.append(local_attn_mask)
712 +
713 + return {
714 + "has_images": True,
715 + "seq_lens": seq_lens,
716 + "global_attn_masks": global_attn_masks,
717 + "local_attn_masks": local_attn_masks,
718 + }
719 +
720 def prepare_attn_masks(
721 self,
722 input_ids: torch.Tensor,
vllm/model_executor/models/siglip.py
+27
@@ -827,6 +827,7 @@ class SiglipVisionModel(nn.Module):
827 ) -> None:
828 super().__init__()
829
830 + self.quant_config = quant_config
831 self.vision_model = SiglipVisionTransformer(
832 config,
833 quant_config,
@@ -911,12 +912,38 @@ class SiglipVisionModel(nn.Module):
912 break
913 else:
914 param = params_dict[name]
915 + param = maybe_swap_ffn_param(
916 + name, param, loaded_weight, params_dict, self.quant_config
917 + )
918 weight_loader = getattr(param, "weight_loader", default_weight_loader)
919 weight_loader(param, loaded_weight)
920 loaded_params.add(name)
921 return loaded_params
922
923
924 +def maybe_swap_ffn_param(
925 + name: str,
926 + param: torch.Tensor,
927 + loaded_weight: torch.Tensor,
928 + params_dict: dict[str, torch.Tensor],
929 + quant_config: QuantizationConfig,
930 +) -> torch.Tensor:
931 + if not (quant_config and quant_config.get_name() == "gguf") or ".fc" not in name:
932 + return param
933 + # Some GGUF models have fc1 and fc2 weights swapped
934 + tp_size = get_tensor_model_parallel_world_size()
935 + output_dim = getattr(param, "output_dim", 0)
936 + output_size = param.size(output_dim) * tp_size
937 + weight_out_size = loaded_weight.size(output_dim)
938 + if ".fc1." in name and output_size != weight_out_size:
939 + new_name = name.replace(".fc1.", ".fc2.")
940 + param = params_dict[new_name]
941 + elif ".fc2." in name and output_size != weight_out_size:
942 + new_name = name.replace(".fc2.", ".fc1.")
943 + param = params_dict[new_name]
944 + return param
945 +
946 +
947 # Adapted from: https://github.com/huggingface/transformers/blob/v4.54.1/src/transformers/models/siglip/modeling_siglip.py#L200
948 class SiglipTextEmbeddings(nn.Module):
949 def __init__(self, config: SiglipTextConfig):
vllm/transformers_utils/config.py
+11
@@ -477,6 +477,17 @@ def is_interleaved(config: PretrainedConfig) -> bool:
477 return False
478
479
480 +def uses_custom_attention_masks(config: PretrainedConfig) -> bool:
481 + """Detect if model uses custom attention mask generation for multimodal.
482 +
483 + Some multimodal models require custom attention masks that enable
484 + bidirectional attention between image tokens while maintaining causal
485 + attention for text tokens. Currently applies to Gemma3 multimodal models.
486 + """
487 + architectures = getattr(config, "architectures", [])
488 + return "Gemma3ForConditionalGeneration" in architectures
489 +
490 +
491 def _maybe_update_auto_config_kwargs(kwargs: dict[str, Any], model_type: str):
492 """
493 Update kwargs for AutoConfig initialization based on model_type
vllm/transformers_utils/gguf_utils.py new
+166
@@ -0,0 +1,166 @@
1 +# SPDX-License-Identifier: Apache-2.0
2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3 +"""GGUF utility functions."""
4 +
5 +from pathlib import Path
6 +
7 +import gguf
8 +from gguf.constants import Keys, VisionProjectorType
9 +from transformers import Gemma3Config, PretrainedConfig, SiglipVisionConfig
10 +
11 +from vllm.logger import init_logger
12 +
13 +logger = init_logger(__name__)
14 +
15 +
16 +def detect_gguf_multimodal(model: str) -> Path | None:
17 + """Check if GGUF model has multimodal projector file.
18 +
19 + Args:
20 + model: Model path string
21 +
22 + Returns:
23 + Path to mmproj file if found, None otherwise
24 + """
25 + if not model.endswith(".gguf"):
26 + return None
27 +
28 + try:
29 + model_path = Path(model)
30 + if not model_path.is_file():
31 + return None
32 +
33 + model_dir = model_path.parent
34 + mmproj_patterns = ["mmproj.gguf", "mmproj-*.gguf", "*mmproj*.gguf"]
35 + for pattern in mmproj_patterns:
36 + mmproj_files = list(model_dir.glob(pattern))
37 + if mmproj_files:
38 + return mmproj_files[0]
39 + return None
40 + except Exception:
41 + return None
42 +
43 +
44 +def extract_vision_config_from_gguf(mmproj_path: str) -> "SiglipVisionConfig | None":
45 + """Extract vision config parameters from mmproj.gguf metadata.
46 +
47 + Reads vision encoder configuration from GGUF metadata fields using
48 + standardized GGUF constants. Automatically detects the projector type
49 + (e.g., gemma3, llama4) and applies model-specific parameters accordingly.
50 +
51 + The function extracts standard CLIP vision parameters from GGUF metadata
52 + and applies projector-type-specific customizations. For unknown projector
53 + types, it uses safe defaults from SiglipVisionConfig.
54 +
55 + Args:
56 + mmproj_path: Path to mmproj.gguf file (str or Path)
57 +
58 + Returns:
59 + SiglipVisionConfig if extraction succeeds, None if any required
60 + field is missing from the GGUF metadata
61 +
62 + Raises:
63 + Exception: Exceptions from GGUF reading (file not found, corrupted
64 + file, etc.) propagate directly from gguf.GGUFReader
65 + """
66 + reader = gguf.GGUFReader(str(mmproj_path))
67 +
68 + # Detect projector type to apply model-specific parameters
69 + projector_type = None
70 + projector_type_field = reader.get_field(Keys.Clip.PROJECTOR_TYPE)
71 + if projector_type_field:
72 + try:
73 + projector_type = bytes(projector_type_field.parts[-1]).decode("utf-8")
74 + except (AttributeError, UnicodeDecodeError) as e:
75 + logger.warning("Failed to decode projector type from GGUF: %s", e)
76 +
77 + # Map GGUF field constants to SiglipVisionConfig parameters.
78 + # Uses official GGUF constants from gguf-py for standardization.
79 + # Format: {gguf_constant: (param_name, dtype)}
80 + VISION_CONFIG_FIELDS = {
81 + Keys.ClipVision.EMBEDDING_LENGTH: ("hidden_size", int),
82 + Keys.ClipVision.FEED_FORWARD_LENGTH: ("intermediate_size", int),
83 + Keys.ClipVision.BLOCK_COUNT: ("num_hidden_layers", int),
84 + Keys.ClipVision.Attention.HEAD_COUNT: ("num_attention_heads", int),
85 + Keys.ClipVision.IMAGE_SIZE: ("image_size", int),
86 + Keys.ClipVision.PATCH_SIZE: ("patch_size", int),
87 + Keys.ClipVision.Attention.LAYERNORM_EPS: ("layer_norm_eps", float),
88 + }
89 +
90 + # Extract and validate all required fields
91 + config_params = {}
92 + for gguf_key, (param_name, dtype) in VISION_CONFIG_FIELDS.items():
93 + field = reader.get_field(gguf_key)
94 + if field is None:
95 + logger.warning(
96 + "Missing required vision config field '%s' in mmproj.gguf",
97 + gguf_key,
98 + )
99 + return None
100 + # Extract scalar value from GGUF field and convert to target type
101 + config_params[param_name] = dtype(field.parts[-1])
102 +
103 + # Apply model-specific parameters based on projector type
104 + if projector_type == VisionProjectorType.GEMMA3:
105 + # Gemma3 doesn't use the vision pooling head (multihead attention)
106 + # This is a vLLM-specific parameter used in SiglipVisionTransformer
107 + config_params["vision_use_head"] = False
108 + logger.info("Detected Gemma3 projector, disabling vision pooling head")
109 + # Add other projector-type-specific customizations here as needed
110 + # elif projector_type == VisionProjectorType.LLAMA4:
111 + # config_params["vision_use_head"] = ...
112 +
113 + # Create config with extracted parameters
114 + # Note: num_channels and attention_dropout use SiglipVisionConfig defaults
115 + # (3 and 0.0 respectively) which are correct for all models
116 + config = SiglipVisionConfig(**config_params)
117 +
118 + if projector_type:
119 + logger.info(
120 + "Extracted vision config from mmproj.gguf (projector_type: %s)",
121 + projector_type,
122 + )
123 + else:
124 + logger.info("Extracted vision config from mmproj.gguf metadata")
125 +
126 + return config
127 +
128 +
129 +def maybe_patch_hf_config_from_gguf(
130 + model: str,
131 + hf_config: PretrainedConfig,
132 +) -> PretrainedConfig:
133 + """Patch HF config for GGUF models.
134 +
135 + Applies GGUF-specific patches to HuggingFace config:
136 + 1. For multimodal models: patches architecture and vision config
137 + 2. For all GGUF models: overrides vocab_size from embedding tensor
138 +
139 + This ensures compatibility with GGUF models that have extended
140 + vocabularies (e.g., Unsloth) where the GGUF file contains more
141 + tokens than the HuggingFace tokenizer config specifies.
142 +
143 + Args:
144 + model: Model path string
145 + hf_config: HuggingFace config to patch in-place
146 +
147 + Returns:
148 + Updated HuggingFace config
149 + """
150 + # Patch multimodal config if mmproj.gguf exists
151 + mmproj_path = detect_gguf_multimodal(model)
152 + if mmproj_path is not None:
153 + vision_config = extract_vision_config_from_gguf(str(mmproj_path))
154 +
155 + # Create HF config for Gemma3 multimodal
156 + text_config = hf_config.get_text_config()
157 + is_gemma3 = hf_config.model_type in ("gemma3", "gemma3_text")
158 + if vision_config is not None and is_gemma3:
159 + new_hf_config = Gemma3Config.from_text_vision_configs(
160 + text_config=text_config,
161 + vision_config=vision_config,
162 + architectures=["Gemma3ForConditionalGeneration"],
163 + )
164 + hf_config = new_hf_config
165 +
166 + return hf_config
vllm/transformers_utils/processor.py
+26 -5
@@ -18,7 +18,7 @@ from transformers.processing_utils import ProcessorMixin
18 from transformers.video_processing_utils import BaseVideoProcessor
19 from typing_extensions import TypeVar
20
21 -from vllm.transformers_utils.utils import convert_model_repo_to_path
21 +from vllm.transformers_utils.utils import check_gguf_file, convert_model_repo_to_path
22 from vllm.utils.func_utils import get_allowed_kwarg_only_overrides
23
24 if TYPE_CHECKING:
@@ -236,9 +236,20 @@ def cached_processor_from_config(
236 processor_cls: type[_P] | tuple[type[_P], ...] = ProcessorMixin,
237 **kwargs: Any,
238 ) -> _P:
239 + if check_gguf_file(model_config.model):
240 + assert not check_gguf_file(model_config.tokenizer), (
241 + "For multimodal GGUF models, the original tokenizer "
242 + "should be used to correctly load processor."
243 + )
244 + model = model_config.tokenizer
245 + revision = model_config.tokenizer_revision
246 + else:
247 + model = model_config.model
248 + revision = model_config.revision
249 +
250 return cached_get_processor_without_dynamic_kwargs(
240 - model_config.model,
241 - revision=model_config.revision,
251 + model,
252 + revision=revision,
253 trust_remote_code=model_config.trust_remote_code,
254 processor_cls=processor_cls, # type: ignore[arg-type]
255 **_merge_mm_kwargs(model_config, processor_cls, **kwargs),
@@ -339,9 +350,19 @@ def cached_image_processor_from_config(
350 model_config: "ModelConfig",
351 **kwargs: Any,
352 ):
353 + if check_gguf_file(model_config.model):
354 + assert not check_gguf_file(model_config.tokenizer), (
355 + "For multimodal GGUF models, the original tokenizer "
356 + "should be used to correctly load image processor."
357 + )
358 + model = model_config.tokenizer
359 + revision = model_config.tokenizer_revision
360 + else:
361 + model = model_config.model
362 + revision = model_config.revision
363 return cached_get_image_processor(
343 - model_config.model,
344 - revision=model_config.revision,
364 + model,
365 + revision=revision,
366 trust_remote_code=model_config.trust_remote_code,
367 **_merge_mm_kwargs(model_config, AutoImageProcessor, **kwargs),
368 )
vllm/transformers_utils/utils.py
+1
@@ -27,6 +27,7 @@ def is_cloud_storage(model_or_path: str) -> bool:
27 return is_s3(model_or_path) or is_gcs(model_or_path)
28
29
30 +@cache
31 def check_gguf_file(model: str | PathLike) -> bool:
32 """Check if the file is a GGUF model."""
33 model = Path(model)
vllm/v1/worker/gpu_model_runner.py
+19
@@ -324,6 +324,7 @@ class GPUModelRunner(
324 # Multi-modal data support
325 self.mm_registry = MULTIMODAL_REGISTRY
326 self.uses_mrope = model_config.uses_mrope
327 + self.uses_custom_attention_masks = model_config.uses_custom_attention_masks
328 self.supports_mm_inputs = self.mm_registry.supports_multimodal_inputs(
329 model_config
330 )
@@ -2346,6 +2347,24 @@ class GPUModelRunner(
2347 **self._init_model_kwargs(num_scheduled_tokens),
2348 **self._extract_mm_kwargs(scheduler_output),
2349 }
2350 +
2351 + # Generate custom attention masks for models that require them.
2352 + # V1 pre-generates embeddings, so forward() skips prepare_attn_masks().
2353 + # Check mm_features (mm_embeds is empty during decode).
2354 + has_mm_features = any(
2355 + req_state.mm_features for req_state in self.requests.values()
2356 + )
2357 + if (
2358 + self.uses_custom_attention_masks
2359 + and has_mm_features
2360 + and hasattr(self.model, "generate_attention_masks")
2361 + ):
2362 + mask_kwargs = self.model.generate_attention_masks(
2363 + self.input_ids.gpu[:num_scheduled_tokens],
2364 + self.positions.gpu[:num_scheduled_tokens],
2365 + mask_dtype=self.model.dtype,
2366 + )
2367 + model_kwargs.update(mask_kwargs)
2368 elif self.enable_prompt_embeds and is_first_rank:
2369 # Get the input embeddings for the tokens that are not input embeds,
2370 # then put them into the appropriate positions.