do not flatten input on linear lora layer

Signed-off-by: NickLucche <[email protected]>

Author: NickLucche

vllm/lora/layers/base_linear.py

@@ -143,7 +143,8 @@ def apply(self,
         # In transformers backend, x and output have extra batch dimension like
         # (1, seq_len, hidden_dim), while punica expects (seq_len, hidden_dim),
         # therefore we need to flatten the batch dimensions.
-        if x.ndim == 3 and output.ndim == 3:
+        if (x.shape[0] ==1 and x.ndim == 3 and
+            output.shape[0] == 1 and output.ndim == 3):
             output = output.flatten(0, 1)
             x = x.flatten(0, 1)
 

vllm/model_executor/models/gemma3n_mm.py

@@ -7,7 +7,6 @@
 import torch
 # yapf: disable
 from torch import nn
-
 from transformers import AutoModel, BatchFeature
 from transformers.models.gemma3n import (Gemma3nAudioConfig,
                                          Gemma3nAudioFeatureExtractor,
@@ -18,6 +17,7 @@
 
 from vllm.config import ModelConfig, SpeechToTextConfig, VllmConfig
 from vllm.config.multimodal import BaseDummyOptions
+from vllm.config.lora import LoRAConfig
 from vllm.inputs.data import PromptType
 from vllm.logger import init_logger
 from vllm.model_executor.layers.layernorm import RMSNorm
@@ -45,8 +45,8 @@
 from vllm.sequence import IntermediateTensors
 from vllm.utils.tensor_schema import TensorSchema, TensorShape
 
-from .interfaces import (MultiModalEmbeddings, SupportsLoRA, SupportsMultiModal,
-                         SupportsTranscription)
+from .interfaces import (MultiModalEmbeddings, SupportsLoRA,
+                         SupportsMultiModal, SupportsTranscription)
 from .utils import (AutoWeightsLoader, WeightsMapper, flatten_bn,
                     init_vllm_registered_model, maybe_prefix)
 
@@ -373,6 +373,7 @@ def __init__(
         self,
         multimodal_config: Union[Gemma3nAudioConfig, Gemma3nVisionConfig],
         text_config: Gemma3nTextConfig,
+        lora_config: Optional[LoRAConfig] = None,
     ):
         super().__init__()
 
@@ -382,9 +383,14 @@ def __init__(
         self.vocab_size = multimodal_config.vocab_size
         self.text_hidden_size = text_config.hidden_size
 
+        lora_vocab = (lora_config.lora_extra_vocab_size *
+                      (lora_config.max_loras or 1)) if lora_config else 0
+        self.vocab_size = self.vocab_size + lora_vocab
+
         self.embedding = VocabParallelEmbedding(
             self.vocab_size,
             self.multimodal_hidden_size,
+            org_num_embeddings=multimodal_config.vocab_size,
         )
 
         self.hard_embedding_norm = RMSNorm(
@@ -427,7 +433,6 @@ def forward(
         if (input_ids is None) ^ (inputs_embeds is not None):
             raise ValueError(
                 "You must specify exactly one of input_ids or inputs_embeds")
-
         if inputs_embeds is not None:
             emb_norm = self.soft_embedding_norm(inputs_embeds)
         else:
@@ -480,13 +485,16 @@ def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):
         self.quant_config = quant_config
         self.multimodal_config = multimodal_config
         self.vocab_size = config.text_config.vocab_size
+        self.lora_config = vllm_config.lora_config
 
         self.vision_tower = AutoModel.from_config(config=config.vision_config)
         self.audio_tower = AutoModel.from_config(config=config.audio_config)
         self.embed_vision = Gemma3nMultimodalEmbedder(config.vision_config,
-                                                      config.text_config)
+                                                      config.text_config,
+                                                      self.lora_config)
         self.embed_audio = Gemma3nMultimodalEmbedder(config.audio_config,
-                                                     config.text_config)
+                                                     config.text_config,
+                                                     self.lora_config)
 
         self.language_model: nn.Module = init_vllm_registered_model(
             vllm_config=vllm_config,
@@ -703,7 +711,7 @@ def get_mm_mapping(self) -> MultiModelKeys:
         return MultiModelKeys.from_string_field(
             language_model="language_model",
             connector="multi_modal_projector",
-            tower_model="vision_tower")
+            tower_model=["vision_tower", "audio_tower"])
 
     @classmethod
     def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]: