@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 5a16fa614c

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 14, 05:30 UTC

[Model] Gemma3n MM (#20495)

Signed-off-by: ShriKode <shrikode@gmail.com> Signed-off-by: NickLucche <nlucches@redhat.com> Signed-off-by: Roger Wang <hey@rogerw.me> Co-authored-by: ShriKode <shrikode@gmail.com> Co-authored-by: Roger Wang <hey@rogerw.me>

Nicolò Lucchesi committed Aug 9, 2025 at 18:56 UTC 5a16fa614c78e1f401125cd7384c602f83cb2160

11 files changed +864 -55

docs/models/supported_models.md

+11 -4

@@ -349,7 +349,7 @@ th {

349 | `GemmaForCausalLM` | Gemma | `google/gemma-2b`, `google/gemma-1.1-2b-it`, etc. | ✅︎ | ✅︎ | ✅︎ |

350 | `Gemma2ForCausalLM` | Gemma 2 | `google/gemma-2-9b`, `google/gemma-2-27b`, etc. | ✅︎ | ✅︎ | ✅︎ |

351 | `Gemma3ForCausalLM` | Gemma 3 | `google/gemma-3-1b-it`, etc. | ✅︎ | ✅︎ | ✅︎ |

353 | `GlmForCausalLM` | GLM-4 | `zai-org/glm-4-9b-chat-hf`, etc. | ✅︎ | ✅︎ | ✅︎ |

354 | `Glm4ForCausalLM` | GLM-4-0414 | `zai-org/GLM-4-32B-0414`, etc. | ✅︎ | ✅︎ | ✅︎ |

355 | `Glm4MoeForCausalLM` | GLM-4.5 | `zai-org/GLM-4.5`, etc. | ✅︎ | ✅︎ | ✅︎ |

@@ -412,9 +412,6 @@ th {

412 !!! note

413 Currently, the ROCm version of vLLM supports Mistral and Mixtral only for context lengths up to 4096.

414

415 -!!! note

416 - Only text inputs are currently supported for `Gemma3nForConditionalGeneration`. To use this model, please upgrade Hugging Face Transformers to version 4.53.0.

417 -

415 ### Pooling Models

416

417 See [this page](./pooling_models.md) for more information on how to use pooling models.

@@ -608,6 +605,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

606 | `FuyuForCausalLM` | Fuyu | T + I | `adept/fuyu-8b`, etc. | | ✅︎ | ✅︎ |

607 | `Gemma3ForConditionalGeneration` | Gemma 3 | T + I+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ | ⚠️ |

608 +| `Gemma3nForConditionalGeneration` | Gemma 3n | T + I + A | `google/gemma-3n-E2B-it`, `google/gemma-3n-E4B-it`, etc. | | | ✅︎ |

609 | `GLM4VForCausalLM`^ | GLM-4V | T + I | `zai-org/glm-4v-9b`, `zai-org/cogagent-9b-20241220`, etc. | ✅︎ | ✅︎ | ✅︎ |

610 | `Glm4vForConditionalGeneration` | GLM-4.1V-Thinking | T + IE+ + VE+ | `zai-org/GLM-4.1V-9B-Thinking`, etc. | ✅︎ | ✅︎ | ✅︎ |

611 | `Glm4vMoeForConditionalGeneration` | GLM-4.5V | T + IE+ + VE+ | `zai-org/GLM-4.5V`, etc. | ✅︎ | ✅︎ | ✅︎ |

@@ -677,6 +675,15 @@ Some models are supported only via the [Transformers backend](#transformers). Th

675

676 This limitation exists because the model's mixed attention pattern (bidirectional for images, causal otherwise) is not yet supported by vLLM's attention backends.

677

678 +!!! note

679 + `Gemma3nForConditionalGeneration` is only supported on V1 due to shared KV caching and it depends on `timm>=1.0.17` to make use of its

680 + MobileNet-v5 vision backbone.

681 +

682 + Performance is not yet fully optimized mainly due to:

683 +

684 + - Both audio and vision MM encoders use `transformers.AutoModel` implementation.

685 + - There's no PLE caching or out-of-memory swapping support, as described in [Google's blog](https://developers.googleblog.com/en/introducing-gemma-3n/). These features might be too model-specific for vLLM, and swapping in particular may be better suited for constrained setups.

686 +

687 !!! note

688 Only `InternVLChatModel` with Qwen2.5 text backbone (`OpenGVLab/InternVL3-2B`, `OpenGVLab/InternVL2.5-1B` etc) has video inputs support currently.

689

examples/offline_inference/audio_language.py

+20

@@ -96,6 +96,25 @@ def run_voxtral(question: str, audio_count: int) -> ModelRequestData:

96 )

99 +# Gemma3N

100 +def run_gemma3n(question: str, audio_count: int) -> ModelRequestData:

101 + model_name = "google/gemma-3n-E2B-it"

102 + engine_args = EngineArgs(

103 + model=model_name,

104 + max_model_len=2048,

105 + max_num_batched_tokens=2048,

106 + max_num_seqs=2,

107 + limit_mm_per_prompt={"audio": audio_count},

108 + enforce_eager=True,

109 + )

110 + prompt = f"<start_of_turn>user\n<audio_soft_token>{question}"

111 + "<end_of_turn>\n<start_of_turn>model\n"

112 + return ModelRequestData(

113 + engine_args=engine_args,

114 + prompt=prompt,

115 + )

116 +

117 +

118 # Granite Speech

119 def run_granite_speech(question: str, audio_count: int) -> ModelRequestData:

120 # NOTE - the setting in this example are somehat different than what is

@@ -331,6 +350,7 @@ def run_whisper(question: str, audio_count: int) -> ModelRequestData:

350

351 model_example_map = {

352 "voxtral": run_voxtral,

353 + "gemma3n": run_gemma3n,

354 "granite_speech": run_granite_speech,

355 "minicpmo": run_minicpmo,

356 "phi4_mm": run_phi4mm,

examples/offline_inference/vision_language.py

+27

@@ -211,7 +211,33 @@ def run_gemma3(questions: list[str], modality: str) -> ModelRequestData:

211 )

212 for question in questions

213 ]

214 + return ModelRequestData(

215 + engine_args=engine_args,

216 + prompts=prompts,

217 + )

218 +

219

220 +# Gemma3N

221 +def run_gemma3n(questions: list[str], modality: str) -> ModelRequestData:

222 + assert modality == "image"

223 + model_name = "google/gemma-3n-E2B-it"

224 +

225 + engine_args = EngineArgs(

226 + model=model_name,

227 + max_model_len=2048,

228 + max_num_seqs=2,

229 + limit_mm_per_prompt={modality: 1},

230 + enforce_eager=True,

231 + )

232 +

233 + prompts = [

234 + (

235 + "<start_of_turn>user\n"

236 + f"<image_soft_token>{question}<end_of_turn>\n"

237 + "<start_of_turn>model\n"

238 + )

239 + for question in questions

240 + ]

241 return ModelRequestData(

242 engine_args=engine_args,

243 prompts=prompts,

@@ -1395,6 +1421,7 @@ model_example_map = {

1421 "florence2": run_florence2,

1422 "fuyu": run_fuyu,

1423 "gemma3": run_gemma3,

1424 + "gemma3n": run_gemma3n,

1425 "glm4v": run_glm4v,

1426 "glm4_1v": run_glm4_1v,

1427 "h2ovl_chat": run_h2ovl,

requirements/test.in

+1 -1

@@ -21,7 +21,7 @@ ray[cgraph,default]>=2.48.0 # Ray Compiled Graph, required by pipeline paralleli

21 sentence-transformers # required for embedding tests

22 soundfile # required for audio tests

23 jiwer # required for audio tests

24 -timm # required for internvl test

24 +timm >=1.0.17 # required for internvl and gemma3n-mm test

25 torch==2.7.1

26 torchaudio==2.7.1

27 torchvision==0.22.1

requirements/test.txt

+1 -1

@@ -1051,7 +1051,7 @@ tiktoken==0.7.0

1051 # via

1052 # lm-eval

1053 # mistral-common

1054 -timm==1.0.15

1054 +timm==1.0.17

1055 # via

1056 # -r requirements/test.in

1057 # open-clip-torch

tests/models/multimodal/processing/test_common.py

+4 -1

@@ -271,6 +271,7 @@ def _test_processing_correctness_one(

271 "microsoft/Florence-2-base",

272 "adept/fuyu-8b",

273 "google/gemma-3-4b-it",

274 + "google/gemma-3n-E2B-it",

275 "zai-org/glm-4v-9b",

276 "zai-org/GLM-4.1V-9B-Thinking",

277 "ibm-granite/granite-speech-3.3-2b",

@@ -315,7 +316,7 @@ def _test_processing_correctness_one(

316 "fixie-ai/ultravox-v0_5-llama-3_2-1b",

317 "openai/whisper-large-v3",

318 "omni-research/Tarsier-7b",

318 - "omni-research/Tarsier2-Recap-7b"

319 + "omni-research/Tarsier2-Recap-7b",

320 ])

321 @pytest.mark.parametrize("hit_rate", [0.3, 0.5, 1.0])

322 @pytest.mark.parametrize("num_batches", [32])

@@ -327,6 +328,8 @@ def test_processing_correctness(

328 num_batches: int,

329 simplify_rate: float,

330 ):

331 + if model_id == "google/gemma-3n-E2B-it":

332 + pytest.skip("Skipping gemma-3n-E2B-it due to transformers #39911 bug.")

333 _test_processing_correctness(

334 model_id,

335 hit_rate=hit_rate,

tests/models/registry.py

+3 -1

@@ -186,7 +186,7 @@ _TEXT_GENERATION_EXAMPLE_MODELS = {

186 "GemmaForCausalLM": _HfExamplesInfo("google/gemma-1.1-2b-it"),

187 "Gemma2ForCausalLM": _HfExamplesInfo("google/gemma-2-9b"),

188 "Gemma3ForCausalLM": _HfExamplesInfo("google/gemma-3-1b-it"),

189 - "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it", # noqa: E501

189 + "Gemma3nForCausalLM": _HfExamplesInfo("google/gemma-3n-E2B-it",

190 min_transformers_version="4.53"),

191 "GlmForCausalLM": _HfExamplesInfo("zai-org/glm-4-9b-chat-hf"),

192 "Glm4ForCausalLM": _HfExamplesInfo("zai-org/GLM-4-9B-0414"),

@@ -391,6 +391,8 @@ _MULTIMODAL_EXAMPLE_MODELS = {

391 "Emu3ForConditionalGeneration": _HfExamplesInfo("BAAI/Emu3-Chat-hf"),

392 "FuyuForCausalLM": _HfExamplesInfo("adept/fuyu-8b"),

393 "Gemma3ForConditionalGeneration": _HfExamplesInfo("google/gemma-3-4b-it"),

394 + "Gemma3nForConditionalGeneration": _HfExamplesInfo("google/gemma-3n-E2B-it", # noqa: E501

395 + min_transformers_version="4.53"),

396 "GraniteSpeechForConditionalGeneration": _HfExamplesInfo("ibm-granite/granite-speech-3.3-2b"), # noqa: E501

397 "GLM4VForCausalLM": _HfExamplesInfo("zai-org/glm-4v-9b",

398 trust_remote_code=True,

tests/test_test.py new

+61

@@ -0,0 +1,61 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +

4 +import pytest

5 +

6 +from vllm import LLM, envs

7 +from vllm.sampling_params import SamplingParams

8 +

9 +if not envs.VLLM_USE_V1:

10 + pytest.skip(

11 + "Skipping V1 tests. Rerun with `VLLM_USE_V1=1` to test.",

12 + allow_module_level=True,

13 + )

14 +

15 +

16 +@pytest.mark.parametrize("model_name", ["Qwen/Qwen2.5-1.5B-Instruct"])

17 +# TODO TPU will appear busy if we fan-out test params here

18 +@pytest.mark.parametrize("n_prompts", [1])

19 +def test_logprobs(model_name: str, n_prompts: int):

20 + """

21 + Request top logprobs with different sampling settings and check

22 + that results contains the requested number, ordered ascendingly.

23 + """

24 +

25 + def check_num_logprobs(logprobs, expected_num: int):

26 + for step in logprobs:

27 + prev_logp = 1.0

28 + # order by rank

29 + sorted_step = dict(

30 + sorted(step.items(), key=lambda item: item[1].rank))

31 +

32 + if len(step) != expected_num:

33 + print("watch out", sorted_step)

34 +

35 + # check results are ordered by prob value

36 + # assert len(step) == expected_num

37 + for rankno, (tid, logp) in enumerate(sorted_step.items()):

38 + assert logp.logprob <= prev_logp

39 + prev_logp = logp.logprob

40 + assert logp.rank == rankno + 1

41 +

42 + llm = LLM(model_name,

43 + enforce_eager=False,

44 + max_num_seqs=1,

45 + max_model_len=128,

46 + max_num_batched_tokens=128)

47 + prompts = [

48 + "Write a short story about a robot that dreams for the first time."

49 + ] * n_prompts

50 + greedy_sampling_params = SamplingParams(temperature=0.0, max_tokens=64,\

51 + logprobs=4)

52 + regular_sampling_params = SamplingParams(temperature=0.4, max_tokens=64,\

53 + logprobs=4)

54 + topkp_sampling_params = SamplingParams(temperature=0.4, max_tokens=64,\

55 + logprobs=4, top_k=12, top_p=0.5)

56 +

57 + for sp in [greedy_sampling_params, regular_sampling_params, \

58 + topkp_sampling_params]:

59 + output = llm.generate(prompts, sp)

60 + for o in output:

61 + check_num_logprobs(o.outputs[0].logprobs, 4)

vllm/model_executor/models/gemma3n.py

+34 -45

@@ -331,14 +331,15 @@ class Gemma3nAttention(nn.Module):

331 config.num_kv_shared_layers)

332 self.is_kv_shared = layer_idx >= first_kv_shared_layer_idx

333

334 + kv_sharing_target_layer_name = None

335 if self.is_kv_shared:

336 # Last full attention layer is 1 before sharing

337 # Last sliding attention layer is 2 before sharing

338 offset = 2 if self.sliding_window is not None else 1

339 kv_shared_layer_index = first_kv_shared_layer_idx - offset

339 - kv_sharing_target_layer_name = f"model.language_model.layers.{kv_shared_layer_index}.self_attn.attn" # noqa: E501

340 - else:

341 - kv_sharing_target_layer_name = None

340 + if kv_shared_layer_index >= 0:

341 + # Only the greater layer is required to specify sharing.

342 + kv_sharing_target_layer_name = f"language_model.model.layers.{kv_shared_layer_index}.self_attn.attn" # noqa: E501

343

344 self.rotary_emb = get_rope(

345 self.head_dim,

@@ -396,6 +397,7 @@ class Gemma3nDecoderLayer(nn.Module):

397 prefix: str = "",

398 ) -> None:

399 super().__init__()

400 + assert isinstance(config, Gemma3nTextConfig)

401 self.altup_active_idx = config.altup_active_idx

402 assert config.altup_correct_scale

403

@@ -537,7 +539,7 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

539

540 def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

541 super().__init__()

540 - config = vllm_config.model_config.hf_config.text_config

542 + config = vllm_config.model_config.hf_config

543 cache_config = vllm_config.cache_config

544 quant_config = vllm_config.quant_config

545 self.config = config

@@ -553,6 +555,7 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

555 config.hidden_size**0.5,

556 dtype=self.embed_tokens.weight.dtype,

557 )

558 + # Additional per-layer embeddings (PLE)

559 self.embed_tokens_per_layer = VocabParallelEmbedding(

560 config.vocab_size_per_layer_input,

561 config.num_hidden_layers * config.hidden_size_per_layer_input,

@@ -636,6 +639,8 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

639 self,

640 input_ids: Optional[torch.Tensor],

641 positions: torch.Tensor,

642 + per_layer_inputs: torch.Tensor,

643 + intermediate_tensors: Optional[IntermediateTensors] = None,

644 inputs_embeds: Optional[torch.Tensor] = None,

645 **kwargs,

646 ) -> Union[torch.Tensor, IntermediateTensors]:

@@ -644,13 +649,6 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

649 else:

650 hidden_states_0 = self.get_input_embeddings(input_ids)

651

647 - # Per layer inputs.

648 - if input_ids is None:

649 - raise ValueError("Passing None for input ids is not supported.")

650 - per_layer_inputs = self.get_per_layer_input_embeddings(input_ids)

651 - per_layer_inputs = per_layer_inputs.reshape(

652 - -1, self.config.num_hidden_layers,

653 - self.config.hidden_size_per_layer_input)

652 per_layer_projection = self.per_layer_model_projection(hidden_states_0)

653 per_layer_projection = per_layer_projection.reshape(

654 *hidden_states_0.shape[:-1],

@@ -659,8 +657,13 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

657 )

658 per_layer_projection = self.per_layer_projection_norm(

659 per_layer_projection)

662 - per_layer_inputs = per_layer_projection + per_layer_inputs

663 - per_layer_inputs *= self.per_layer_input_scale

660 +

661 + if per_layer_inputs is not None:

662 + # Profiling run does not compute per_layer_inputs

663 + per_layer_inputs = per_layer_projection + per_layer_inputs

664 + per_layer_inputs *= self.per_layer_input_scale

665 + else:

666 + per_layer_inputs = per_layer_projection

667

668 # Altup embed.

669 hidden_states = [hidden_states_0] * self.config.altup_num_inputs

@@ -760,29 +763,7 @@ class Gemma3nTextModel(nn.Module, SupportsQuant):

763 return loaded_params

764

765

763 -class Gemma3nModel(nn.Module):

764 -

765 - def __init__(self, vllm_config: VllmConfig, prefix: str = ""):

766 - super().__init__()

767 - self.language_model = Gemma3nTextModel(vllm_config=vllm_config,

768 - prefix=maybe_prefix(

769 - prefix, "language_model"))

770 -

771 - def forward(

772 - self,

773 - input_ids: Optional[torch.Tensor],

774 - positions: torch.Tensor,

775 - intermediate_tensors: Optional[IntermediateTensors] = None,

776 - inputs_embeds: Optional[torch.Tensor] = None,

777 - **kwargs,

778 - ) -> torch.Tensor:

779 - return self.language_model(input_ids=input_ids,

780 - positions=positions,

781 - inputs_embeds=inputs_embeds,

782 - **kwargs)

783 -

784 -

785 -class Gemma3nForConditionalGeneration(nn.Module, SupportsQuant):

766 +class Gemma3nForCausalLM(nn.Module):

767 packed_modules_mapping = {

768 "qkv_proj": [

769 "q_proj",

@@ -802,25 +783,33 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsQuant):

783 super().__init__()

784 self.config = config

785 self.cache_config = vllm_config.cache_config

805 - self.model = Gemma3nModel(vllm_config=vllm_config,

806 - prefix=maybe_prefix(prefix, "model"))

786 + self.model = Gemma3nTextModel(vllm_config=vllm_config,

787 + prefix=maybe_prefix(prefix, "model"))

788 self.logits_processor = LogitsProcessor(

808 - config.text_config.vocab_size,

809 - soft_cap=config.text_config.final_logit_softcapping)

789 + config.vocab_size, soft_cap=config.final_logit_softcapping)

790

791 def get_input_embeddings(self, input_ids: torch.Tensor) -> torch.Tensor:

812 - return self.model.language_model.get_input_embeddings(input_ids)

792 + return self.model.get_input_embeddings(input_ids)

793

794 def forward(

795 self,

796 input_ids: torch.Tensor,

797 positions: torch.Tensor,

798 + *,

799 + per_layer_inputs: Optional[torch.Tensor] = None,

800 intermediate_tensors: Optional[IntermediateTensors] = None,

801 inputs_embeds: Optional[torch.Tensor] = None,

802 **kwargs,

803 ) -> Union[torch.Tensor, IntermediateTensors]:

822 - hidden_states = self.model(input_ids, positions, intermediate_tensors,

823 - inputs_embeds, **kwargs)

804 +

805 + hidden_states = self.model(

806 + input_ids,

807 + positions,

808 + per_layer_inputs=per_layer_inputs,

809 + intermediate_tensors=intermediate_tensors,

810 + inputs_embeds=inputs_embeds,

811 + **kwargs,

812 + )

813 return hidden_states

814

815 def compute_logits(

@@ -828,8 +817,8 @@ class Gemma3nForConditionalGeneration(nn.Module, SupportsQuant):

817 hidden_states: torch.Tensor,

818 sampling_metadata: Optional[SamplingMetadata],

819 ) -> Optional[torch.Tensor]:

831 - logits = self.logits_processor(self.model.language_model.embed_tokens,

832 - hidden_states, sampling_metadata)

820 + logits = self.logits_processor(self.model.embed_tokens, hidden_states,

821 + sampling_metadata)

822 return logits

823

824 def load_weights(self, weights: Iterable[tuple[str,

vllm/model_executor/models/gemma3n_mm.py new

+700

@@ -0,0 +1,700 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

3 +from collections.abc import Iterable, Mapping, Sequence

4 +from typing import Any, Optional, TypedDict, Union, cast

5 +

6 +import torch

7 +from torch import nn

8 +from transformers import AutoModel, BatchFeature

9 +from transformers.models.gemma3n import (Gemma3nAudioConfig,

10 + Gemma3nAudioFeatureExtractor,

11 + Gemma3nConfig, Gemma3nProcessor,

12 + Gemma3nTextConfig,

13 + Gemma3nVisionConfig)

14 +from transformers.models.siglip import SiglipImageProcessorFast

15 +

16 +from vllm.config import VllmConfig

17 +from vllm.logger import init_logger

18 +from vllm.model_executor.layers.layernorm import RMSNorm

19 +from vllm.model_executor.layers.linear import RowParallelLinear

20 +from vllm.model_executor.layers.vocab_parallel_embedding import (

21 + VocabParallelEmbedding)

22 +from vllm.model_executor.models.gemma3n import Gemma3nForCausalLM

23 +from vllm.model_executor.models.module_mapping import MultiModelKeys

24 +from vllm.model_executor.sampling_metadata import SamplingMetadata

25 +from vllm.multimodal import MULTIMODAL_REGISTRY

26 +from vllm.multimodal.inputs import (MultiModalDataDict, MultiModalFieldConfig,

27 + MultiModalKwargs)

28 +from vllm.multimodal.parse import (ImageProcessorItems, MultiModalDataItems,

29 + MultiModalDataParser)

30 +# yapf: disable

31 +from vllm.multimodal.processing import (BaseMultiModalProcessor,

32 + BaseProcessingInfo, BoundPromptUpdate,

33 + PlaceholderFeaturesInfo,

34 + PromptReplacement, PromptTargetMatch,

35 + PromptUpdate, PromptUpdateDetails,

36 + find_mm_placeholders,

37 + replace_token_matches)

38 +# yapf: enable

39 +from vllm.multimodal.profiling import BaseDummyInputsBuilder

40 +from vllm.sequence import IntermediateTensors

41 +

42 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal

43 +from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,

44 + init_vllm_registered_model, maybe_prefix,

45 + merge_multimodal_embeddings)

46 +

47 +logger = init_logger(__name__)

48 +

49 +# This should be based on model config but we hardcode them for now.

50 +TOKENS_PER_IMAGE = 256

51 +TOKENS_PER_AUDIO = 188

52 +

53 +

54 +class Gemma3nImagePixelInputs(TypedDict):

55 + pixel_values: torch.Tensor

56 + """Shape: `(batch_size * num_images, num_channels, height, width)`"""

57 +

58 +

59 +class Gemma3nAudioInputs(TypedDict):

60 + input_features: torch.Tensor

61 + """Shape: `(batch_size * num_audio, seq_length, num_features)`"""

62 + input_features_mask: torch.Tensor

63 + """Shape: `(batch_size * num_audio, seq_length)`"""

64 +

65 +

66 +Gemma3nImageInputs = Gemma3nImagePixelInputs

67 +

68 +

69 +class Gemma3nProcessingInfo(BaseProcessingInfo):

70 +

71 + def get_hf_config(self):

72 + return self.ctx.get_hf_config(Gemma3nConfig)

73 +

74 + def get_hf_processor(self, **kwargs: object):

75 + return self.ctx.get_hf_processor(Gemma3nProcessor, **kwargs)

76 +

77 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

78 + return {"image": None, "audio": None}

79 +

80 + def get_max_tokens_per_item(

81 + self, seq_len: int,

82 + mm_counts: Mapping[str, int]) -> Optional[Mapping[str, int]]:

83 +

84 + return {"image": TOKENS_PER_IMAGE, "audio": TOKENS_PER_AUDIO}

85 +

86 + def get_image_repl(

87 + self,

88 + *,

89 + image_width: int,

90 + image_height: int,

91 + processor: Optional[Gemma3nProcessor],

92 + ) -> str:

93 + """

94 + Get the replacement text for image tokens.

95 +

96 + For Gemma3n, this should return the full_image_sequence which includes

97 + BOI token, repeated image tokens, and EOI token.

98 + """

99 + if processor is None:

100 + processor = self.get_hf_processor()

101 +

102 + return PromptUpdateDetails.select_token_id(

103 + processor.full_image_sequence, processor.image_token_id)

104 +

105 + def get_audio_repl(

106 + self,

107 + *,

108 + processor: Optional[Gemma3nProcessor],

109 + ) -> str:

110 + """

111 + Get the replacement text for audio tokens.

112 +

113 + For Gemma3n, this should return the full_audio_sequence which includes

114 + BOA token, repeated audio tokens, and EOA token.

115 + """

116 + if processor is None:

117 + processor = self.get_hf_processor()

118 +

119 + # Return the full audio sequence as defined by the processor

120 + return PromptUpdateDetails.select_token_id(

121 + processor.full_audio_sequence, processor.audio_token_id)

122 +

123 +

124 +class Gemma3nDummyInputsBuilder(BaseDummyInputsBuilder[Gemma3nProcessingInfo]):

125 +

126 + def get_dummy_text(self, mm_counts: Mapping[str, int]) -> str:

127 + num_images = mm_counts.get("image", 0)

128 + num_audios = mm_counts.get("audio", 0)

129 +

130 + processor = self.info.get_hf_processor()

131 + image_token = processor.image_token

132 + audio_token = processor.audio_token

133 +

134 + return image_token * num_images + audio_token * num_audios

135 +

136 + def get_dummy_mm_data(

137 + self,

138 + seq_len: int,

139 + mm_counts: Mapping[str, int],

140 + ) -> MultiModalDataDict:

141 + num_images = mm_counts.get("image", 0)

142 + num_audios = mm_counts.get("audio", 0)

143 + processor = self.info.get_hf_processor()

144 + audio_feature_extractor: Gemma3nAudioFeatureExtractor = processor.feature_extractor # noqa: E501

145 + audio_len = audio_feature_extractor.fft_length

146 + image_processor: SiglipImageProcessorFast = processor.image_processor

147 + img_width = image_processor.size.get("width", 224)

148 + img_height = image_processor.size.get("height", 224)

149 +

150 + return {

151 + "image":

152 + self._get_dummy_images(width=img_width,

153 + height=img_height,

154 + num_images=num_images),

155 + "audio":

156 + self._get_dummy_audios(length=audio_len, num_audios=num_audios)

157 + }

158 +

159 +

160 +class Gemma3nMultiModalProcessor(BaseMultiModalProcessor[Gemma3nProcessingInfo]

161 + ):

162 +

163 + def _get_data_parser(self) -> MultiModalDataParser:

164 + feature_extractor = self.info.get_hf_processor().feature_extractor

165 + return MultiModalDataParser(target_sr=feature_extractor.sampling_rate)

166 +

167 + def _call_hf_processor(

168 + self,

169 + prompt: str,

170 + mm_data: Mapping[str, object],

171 + mm_kwargs: Mapping[str, object],

172 + tok_kwargs: Mapping[str, object],

173 + ) -> BatchFeature:

174 +

175 + # HF Transformers audio processor no longer accepts `audios` key.

176 + # We pop `audios` and replace it with `audio` key to surpress

177 + # the warning.

178 + if 'audios' in mm_data:

179 + mm_data['audio'] = mm_data.pop('audios')

180 + processed_outputs = super()._call_hf_processor(

181 + prompt,

182 + mm_data,

183 + mm_kwargs,

184 + tok_kwargs,

185 + )

186 + if 'input_features' in processed_outputs:

187 + # Avoid padding since we need the output of each item to be

188 + # independent of other items for the cache to work correctly

189 + unpadded_features = [

190 + f[mask] for f, mask in zip(

191 + processed_outputs["input_features"],

192 + processed_outputs["input_features_mask"],

193 + )

194 + ]

195 + processed_outputs["input_features"] = unpadded_features

196 + return processed_outputs

197 +

198 + def _get_mm_fields_config(

199 + self,

200 + hf_inputs: BatchFeature,

201 + hf_processor_mm_kwargs: Mapping[str, object],

202 + ) -> Mapping[str, MultiModalFieldConfig]:

203 +

204 + return dict(pixel_values=MultiModalFieldConfig.batched("image"),

205 + input_features=MultiModalFieldConfig.batched("audio"),

206 + input_features_mask=MultiModalFieldConfig.batched("audio"))

207 +

208 + def _get_prompt_updates(

209 + self,

210 + mm_items: MultiModalDataItems,

211 + hf_processor_mm_kwargs: Mapping[str, Any],

212 + out_mm_kwargs: MultiModalKwargs,

213 + ) -> Sequence[PromptUpdate]:

214 + hf_processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

215 +

216 + prompt_updates = []

217 +

218 + # Handle image tokens

219 + if "image" in mm_items:

220 + image_token = hf_processor.image_token

221 +

222 + def get_replacement_image(item_idx: int):

223 + images = mm_items.get_items("image", ImageProcessorItems)

224 + image_size = images.get_image_size(item_idx)

225 + return self.info.get_image_repl(

226 + image_width=image_size.width,

227 + image_height=image_size.height,

228 + processor=hf_processor,

229 + )

230 +

231 + prompt_updates.append(

232 + PromptReplacement(

233 + modality="image",

234 + target=image_token,

235 + replacement=get_replacement_image,

236 + ))

237 +

238 + # Handle audio tokens

239 + if "audio" in mm_items:

240 + audio_token = hf_processor.audio_token

241 +

242 + def get_replacement_audio(item_idx: int):

243 + return self.info.get_audio_repl(processor=hf_processor, )

244 +

245 + prompt_updates.append(

246 + PromptReplacement(

247 + modality="audio",

248 + target=audio_token,

249 + replacement=get_replacement_audio,

250 + ))

251 +

252 + return prompt_updates

253 +

254 + def _apply_token_matches(

255 + self,

256 + prompt: list[int],

257 + mm_matches: Mapping[str, Sequence[PromptTargetMatch]],

258 + mm_item_counts: Mapping[str, int],

259 + ) -> list[int]:

260 + token_ids = super()._apply_token_matches(

261 + prompt,

262 + mm_matches,

263 + mm_item_counts,

264 + )

265 +

266 + # "\n\n\n" and "\n\n\n\n" are single tokens

267 + # Since our replacement can insert "\n\n" next to "\n"

268 + # tokens, we have to combine them to be consistent with

269 + # the output of the tokenizer

270 + tokenizer = self.info.get_tokenizer()

271 + vocab = tokenizer.get_vocab()

272 + newline_1 = vocab["\n"]

273 + newline_2 = vocab["\n\n"]

274 + newline_3 = vocab["\n\n\n"]

275 + newline_4 = vocab["\n\n\n\n"]

276 +

277 + token_ids = replace_token_matches(

278 + token_ids,

279 + [newline_1, newline_2],

280 + [newline_3],

281 + )

282 + token_ids = replace_token_matches(

283 + token_ids,

284 + [newline_2, newline_1],

285 + [newline_3],

286 + )

287 + token_ids = replace_token_matches(

288 + token_ids,

289 + [newline_2, newline_2],

290 + [newline_4],

291 + )

292 +

293 + return token_ids

294 +

295 + def _find_mm_placeholders(

296 + self,

297 + mm_prompt_updates: Mapping[str, Sequence[BoundPromptUpdate]],

298 + new_token_ids: list[int],

299 + mm_item_counts: Mapping[str, int],

300 + ) -> Mapping[str, list[PlaceholderFeaturesInfo]]:

301 + # We need to detect "\n\n" inside "\n\n\n" and "\n\n\n\n"

302 + tokenizer = self.info.get_tokenizer()

303 + vocab = tokenizer.get_vocab()

304 + newline_1 = vocab["\n"]

305 + newline_2 = vocab["\n\n"]

306 + newline_3 = vocab["\n\n\n"]

307 + newline_4 = vocab["\n\n\n\n"]

308 +

309 + def get_repl_toks(tok: int) -> list[int]:

310 + if tok == newline_3:

311 + return [newline_1, newline_2]

312 + if tok == newline_4:

313 + return [newline_2, newline_2]

314 +

315 + return [tok]

316 +

317 + repl_token_ids = list[int]()

318 + repl_orig_idxs = list[int]()

319 + for orig_idx, orig_tok in enumerate(new_token_ids):

320 + repl_toks = get_repl_toks(orig_tok)

321 + repl_token_ids.extend(repl_toks)

322 + repl_orig_idxs.extend(orig_idx for _ in range(len(repl_toks)))

323 +

324 + repls = find_mm_placeholders(mm_prompt_updates, repl_token_ids,

325 + mm_item_counts)

326 +

327 + return {

328 + modality: [

329 + PlaceholderFeaturesInfo(

330 + modality=p.modality,

331 + item_idx=p.item_idx,

332 + start_idx=repl_orig_idxs[p.start_idx],

333 + tokens=p.tokens,

334 + is_embed=p.is_embed,

335 + ) for p in placeholders

336 + ]

337 + for modality, placeholders in repls.items()

338 + }

339 +

340 +

341 +class Gemma3nMultimodalEmbedder(nn.Module):

342 + """Embeds token ids or soft tokens for multimodal content into language

343 + model space."""

344 +

345 + def __init__(

346 + self,

347 + multimodal_config: Union[Gemma3nAudioConfig, Gemma3nVisionConfig],

348 + text_config: Gemma3nTextConfig,

349 + ):

350 + super().__init__()

351 +

352 + self.multimodal_hidden_size = multimodal_config.hidden_size

353 + self.eps = multimodal_config.rms_norm_eps

354 + self.vocab_offset = multimodal_config.vocab_offset

355 + self.vocab_size = multimodal_config.vocab_size

356 + self.text_hidden_size = text_config.hidden_size

357 +

358 + self.embedding = VocabParallelEmbedding(

359 + self.vocab_size,

360 + self.multimodal_hidden_size,

361 + )

362 +

363 + self.hard_embedding_norm = RMSNorm(

364 + self.multimodal_hidden_size,

365 + eps=self.eps,

366 + )

367 +

368 + self.soft_embedding_norm = RMSNorm(

369 + self.multimodal_hidden_size,

370 + eps=self.eps,

371 + )

372 +

373 + self.embedding_projection = RowParallelLinear(

374 + self.multimodal_hidden_size,

375 + self.text_hidden_size,

376 + bias=False,

377 + )

378 +

379 + self.embedding_post_projection_norm = RMSNorm(

380 + self.text_hidden_size,

381 + eps=self.eps,

382 + has_weight=False,

383 + )

384 +

385 + def forward(

386 + self,

387 + input_ids: Optional[torch.LongTensor] = None,

388 + inputs_embeds: Optional[torch.Tensor] = None,

389 + ) -> torch.Tensor:

390 + """Embeds token ids or soft tokens for multimodal content into language model space.

391 +

392 + Args:

393 + input_ids: A torch.LongTensor containing the token ids to embed. Values should be in the range

394 + `[vocab_offset, vocab_offset + vocab_size)`.

395 + inputs_embeds: A torch.Tensor containing the soft tokens to embed.

396 +

397 + Returns:

398 + A torch.Tensor of embeddings with shape `[batch_size, seq_len, self.config.text_config.hidden_size]`.

399 + """ # noqa: E501

400 + if (input_ids is None) ^ (inputs_embeds is not None):

401 + raise ValueError(

402 + "You must specify exactly one of input_ids or inputs_embeds")

403 +

404 + if inputs_embeds is not None:

405 + emb_norm = self.soft_embedding_norm(inputs_embeds)

406 + else:

407 + hard_emb = self.embedding(input_ids - self.vocab_offset)

408 + emb_norm = self.hard_embedding_norm(hard_emb)

409 +

410 + emb_norm_proj, _ = self.embedding_projection(emb_norm)

411 + return self.embedding_post_projection_norm(emb_norm_proj)

412 +

413 +

414 +@MULTIMODAL_REGISTRY.register_processor(Gemma3nMultiModalProcessor,

415 + info=Gemma3nProcessingInfo,

416 + dummy_inputs=Gemma3nDummyInputsBuilder)

417 +class Gemma3nForConditionalGeneration(nn.Module, SupportsMultiModal):

418 + packed_modules_mapping = {

419 + "qkv_proj": [

420 + "q_proj",

421 + "k_proj",

422 + "v_proj",

423 + ],

424 + "gate_up_proj": [

425 + "gate_proj",

426 + "up_proj",

427 + ],

428 + }

429 +

430 + hf_to_vllm_mapper = WeightsMapper(

431 + orig_to_new_prefix={

432 + # mapping for new names in checkpoint saved after transformers v4.52

433 + "model.embed_audio.": "embed_audio.",

434 + "model.embed_vision.": "embed_vision.",

435 + "model.language_model.": "language_model.model.",

436 + "model.vision_tower.": "vision_tower.",

437 + "model.audio_tower.": "audio_tower.",

438 + "model.multi_modal_projector.": "multi_modal_projector.",

439 + "lm_head.": "language_model.lm_head.",

440 + "model": "language_model.model",

441 + })

442 +

443 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

444 + super().__init__()

445 + config = vllm_config.model_config.hf_config

446 + quant_config = vllm_config.quant_config

447 + multimodal_config = vllm_config.model_config.multimodal_config

448 + self.config = config

449 + self.quant_config = quant_config

450 + self.multimodal_config = multimodal_config

451 + self.vocab_size = config.text_config.vocab_size

452 +

453 + self.sliding_window = getattr(config.text_config,

454 + "interleaved_sliding_window", None)

455 +

456 + self.vision_tower = AutoModel.from_config(config=config.vision_config)

457 + self.audio_tower = AutoModel.from_config(config=config.audio_config)

458 + self.embed_vision = Gemma3nMultimodalEmbedder(config.vision_config,

459 + config.text_config)

460 + self.embed_audio = Gemma3nMultimodalEmbedder(config.audio_config,

461 + config.text_config)

462 +

463 + self.language_model: nn.Module = init_vllm_registered_model(

464 + vllm_config=vllm_config,

465 + hf_config=config.text_config,

466 + prefix=maybe_prefix(prefix, "language_model"),

467 + architectures=["Gemma3nForCausalLM"],

468 + )

469 + self.language_model = cast(Gemma3nForCausalLM, self.language_model)

470 + # NOTE (NickLucche) In order to be compatible with cudagraph, the

471 + # buffer needs to be consistent, so we pre-allocate here.

472 + self.per_layer_embeddings = torch.zeros(

473 + vllm_config.scheduler_config.max_num_batched_tokens,

474 + self.config.text_config.num_hidden_layers,

475 + self.config.text_config.hidden_size_per_layer_input,

476 + device=self.language_model.model.embed_tokens.weight.device,

477 + dtype=self.language_model.model.embed_tokens.weight.dtype)

478 +

479 + @property

480 + def dtype(self):

481 + return next(self.parameters()).dtype

482 +

483 + def _validate_pixel_values(self, data: torch.Tensor) -> torch.Tensor:

484 + # TODO check if there are any

485 + return data

486 +

487 + def _parse_and_validate_image_input(

488 + self, **kwargs: object) -> Optional[Gemma3nImageInputs]:

489 + pixel_values = kwargs.pop("pixel_values", None)

490 + image_embeds = kwargs.pop("image_embeds", None)

491 + # TODO is this the case?

492 + assert image_embeds is None, "Gemma3n does not support image_embeds."

493 + if pixel_values is None:

494 + return None

495 +

496 + if not isinstance(pixel_values, (torch.Tensor, list)):

497 + raise ValueError("Incorrect type of pixel values. "

498 + f"Got type: {type(pixel_values)}")

499 +

500 + pixel_values = flatten_bn(pixel_values, concat=True)

501 + pixel_values = pixel_values.contiguous()

502 +

503 + return Gemma3nImagePixelInputs(

504 + pixel_values=self._validate_pixel_values(pixel_values), )

505 +

506 + def _parse_and_validate_audio_input(

507 + self, **kwargs: object) -> Optional[Gemma3nAudioInputs]:

508 + input_features = kwargs.pop("input_features", None)

509 + if input_features is None:

510 + return None

511 +

512 + input_features_mask = kwargs.pop("input_features_mask", None)

513 + if input_features_mask is None:

514 + return None

515 +

516 + return Gemma3nAudioInputs(

517 + input_features=input_features,

518 + input_features_mask=input_features_mask,

519 + )

520 +

521 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

522 + mm_input_by_modality = {}

523 +

524 + # Preserve the order of modalities if there are multiple of them

525 + # from the order of kwargs.

526 + for input_key in kwargs:

527 + if input_key in ("pixel_values", "image_embeds"

528 + ) and "image" not in mm_input_by_modality:

529 + mm_input_by_modality[

530 + "image"] = self._parse_and_validate_image_input(**kwargs)

531 + if input_key == "input_features" \

532 + and "audio" not in mm_input_by_modality:

533 + mm_input_by_modality[

534 + "audio"] = self._parse_and_validate_audio_input(**kwargs)

535 + return mm_input_by_modality

536 +

537 + def _process_image_input(

538 + self,

539 + image_input: Gemma3nImageInputs,

540 + ) -> list[torch.Tensor]:

541 + assert self.vision_tower is not None

542 +

543 + pixel_values = image_input["pixel_values"]

544 + vision_outputs = self.vision_tower(pixel_values=pixel_values,

545 + do_pooling=False,

546 + return_dict=True).last_hidden_state

547 + # TODO try to avoid copy here

548 + # (batch, channels, height, width) to (batch, height * width, channels)

549 + vision_outputs = vision_outputs.reshape(

550 + vision_outputs.shape[0],

551 + self.config.vision_config.hidden_size,

552 + self.config.vision_soft_tokens_per_image,

553 + ).permute(0, 2, 1).contiguous()

554 + # Normalize and embed the soft tokens into language model space.

555 + vision_outputs *= self.config.vision_config.hidden_size**0.5

556 + # Return a list of embeddings instead of a batched tensor

557 + return self.embed_vision(inputs_embeds=vision_outputs).unbind(0)

558 +

559 + def _process_audio_input(

560 + self,

561 + audio_input: Gemma3nAudioInputs,

562 + ) -> list[torch.Tensor]:

563 + assert self.audio_tower is not None

564 + input_features = audio_input["input_features"].squeeze(1)

565 + input_features_mask = audio_input["input_features_mask"].squeeze(1)

566 + audio_outputs, audio_mask = self.audio_tower(input_features,

567 + ~input_features_mask)

568 + audio_features = self.embed_audio(inputs_embeds=audio_outputs)

569 +

570 + # ruff: noqa

571 + # The Gemma3nProcessor expects all audio will be 30s in length and inserts 188 audio soft tokens into the

572 + # text to account for this. However, the audio preprocessing and encoder do not gurarantee they will

573 + # produce 188 soft tokens; they will produce at most that many tokens, but they may produce fewer tokens

574 + # depending on the length of the longest audio input in the batch. When we encounter this situation, we pad

575 + # the audio feature out to 188 soft tokens with the emebedding of the last token in the embed_audio vocab.

576 + # TODO precompute and cache padding

577 + audio_padding_toks = torch.tensor([[self.vocab_size - 1]],

578 + dtype=torch.long,

579 + device=audio_features.device)

580 + audio_padding_embs = self.embed_audio(input_ids=audio_padding_toks)

581 + audio_features = torch.where(audio_mask.unsqueeze(-1),

582 + audio_padding_embs, audio_features)

583 +

584 + audio_batch_size, audio_seq_len, audio_embed_dim = audio_features.shape

585 + extra_padding_tokens = self.config.audio_soft_tokens_per_image - audio_seq_len # noqa: E501

586 + extra_padding_features = audio_padding_embs.expand(

587 + audio_batch_size, extra_padding_tokens, audio_embed_dim)

588 +

589 + audio_features = torch.cat((audio_features, extra_padding_features),

590 + dim=1)

591 + # Return a list of embeddings instead of a batched tensor

592 + return audio_features.unbind(0)

593 +

594 + def get_language_model(self) -> torch.nn.Module:

595 + return self.language_model

596 +

597 + def get_multimodal_embeddings(self,

598 + **kwargs: object) -> MultiModalEmbeddings:

599 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(

600 + **kwargs)

601 + if mm_input_by_modality is None:

602 + return []

603 +

604 + multimodal_embeddings: list[torch.Tensor] = []

605 +

606 + # NOTE: It is important to iterate over the keys in this dictionary

607 + # to preserve the order of the modalities.

608 + for modality in mm_input_by_modality:

609 + multimodal_input = mm_input_by_modality[modality]

610 + if modality == "image":

611 + vision_embeddings = self._process_image_input(multimodal_input)

612 + multimodal_embeddings.extend(vision_embeddings)

613 + if modality == "audio":

614 + audio_embeddings = self._process_audio_input(multimodal_input)

615 + multimodal_embeddings.extend(audio_embeddings)

616 + return multimodal_embeddings

617 +

618 + def get_input_embeddings(

619 + self,

620 + input_ids: torch.Tensor,

621 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

622 + ) -> torch.Tensor:

623 + inputs_embeds = self.language_model.get_input_embeddings(input_ids)

624 + # NOTE (NickLucche) Each pass needs tokens to compute PLE so we cache

625 + # them here, as the model forward has only access to the input_embeds.

626 + if input_ids is not None:

627 + per_layer_inputs = self.language_model.model.get_per_layer_input_embeddings(

628 + input_ids)

629 + per_layer_inputs = per_layer_inputs.reshape(

630 + -1, self.config.text_config.num_hidden_layers,

631 + self.config.text_config.hidden_size_per_layer_input)

632 + self.per_layer_embeddings[:per_layer_inputs.shape[0]].copy_(

633 + per_layer_inputs)

634 +

635 + if multimodal_embeddings is not None \

636 + and len(multimodal_embeddings) != 0:

637 + inputs_embeds = merge_multimodal_embeddings(

638 + input_ids,

639 + inputs_embeds,

640 + multimodal_embeddings,

641 + # NOTE: this order of processing mm items is important

642 + [self.config.image_token_id, self.config.audio_token_id])

643 + return inputs_embeds

644 +

645 + def forward(self,

646 + input_ids: torch.Tensor,

647 + positions: torch.Tensor,

648 + intermediate_tensors: Optional[IntermediateTensors] = None,

649 + inputs_embeds: Optional[torch.Tensor] = None,

650 + **kwargs: object) -> IntermediateTensors:

651 + if intermediate_tensors is not None:

652 + inputs_embeds = None

653 +

654 + # NOTE (NickLucche) During profiling, `get_input_embeddings` is not

655 + # called, hence we don't have input_ids to compute PLEs. We simply

656 + # select a chunk of pre-allocated PLEs. During normal execution,

657 + # `get_input_embeddings` is called before forward, hence this slice

658 + # will contain PLEs computed from the actual input_ids.

659 + per_layer_inputs = self.per_layer_embeddings[:inputs_embeds.shape[0]]

660 +

661 + hidden_states = self.language_model.model(

662 + input_ids,

663 + positions,

664 + per_layer_inputs=per_layer_inputs,

665 + intermediate_tensors=intermediate_tensors,

666 + inputs_embeds=inputs_embeds,

667 + **kwargs)

668 +

669 + return hidden_states

670 +

671 + def compute_logits(

672 + self,

673 + hidden_states: torch.Tensor,

674 + sampling_metadata: SamplingMetadata,

675 + ) -> Optional[torch.Tensor]:

676 + return self.language_model.compute_logits(hidden_states,

677 + sampling_metadata)

678 +

679 + def load_weights(self, weights: Iterable[tuple[str,

680 + torch.Tensor]]) -> set[str]:

681 + loader = AutoWeightsLoader(self)

682 + return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

683 +

684 + def get_mm_mapping(self) -> MultiModelKeys:

685 + """

686 + Get the module prefix in multimodal models

687 + """

688 + return MultiModelKeys.from_string_field(

689 + language_model="language_model",

690 + connector="multi_modal_projector",

691 + tower_model="vision_tower")

692 +

693 + @classmethod

694 + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

695 + if modality == "image":

696 + return "<image_soft_token>"

697 + elif modality == "audio":

698 + return "<audio_soft_token>"

699 + else:

700 + raise ValueError(f"Unsupported modality: {modality}")

vllm/model_executor/models/registry.py

+2 -2

@@ -69,8 +69,7 @@ _TEXT_GENERATION_MODELS = {

69 "GemmaForCausalLM": ("gemma", "GemmaForCausalLM"),

70 "Gemma2ForCausalLM": ("gemma2", "Gemma2ForCausalLM"),

71 "Gemma3ForCausalLM": ("gemma3", "Gemma3ForCausalLM"),

72 - #TODO(ywang96): Support multimodal gemma3n

73 - "Gemma3nForConditionalGeneration": ("gemma3n", "Gemma3nForConditionalGeneration"), # noqa: E501

72 + "Gemma3nForCausalLM": ("gemma3n", "Gemma3nForCausalLM"),

73 "GlmForCausalLM": ("glm", "GlmForCausalLM"),

74 "Glm4ForCausalLM": ("glm4", "Glm4ForCausalLM"),

75 "Glm4MoeForCausalLM": ("glm4_moe", "Glm4MoeForCausalLM"),

@@ -205,6 +204,7 @@ _MULTIMODAL_MODELS = {

204 "DeepseekVLV2ForCausalLM": ("deepseek_vl2", "DeepseekVLV2ForCausalLM"),

205 "FuyuForCausalLM": ("fuyu", "FuyuForCausalLM"),

206 "Gemma3ForConditionalGeneration": ("gemma3_mm", "Gemma3ForConditionalGeneration"), # noqa: E501

207 + "Gemma3nForConditionalGeneration": ("gemma3n_mm", "Gemma3nForConditionalGeneration"), # noqa: E501

208 "GLM4VForCausalLM": ("glm4v", "GLM4VForCausalLM"),

209 "Glm4vForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), # noqa: E501

210 "Glm4vMoeForConditionalGeneration": ("glm4_1v", "Glm4vForConditionalGeneration"), # noqa: E501