@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 95c9144424

[LoRA][Gemma4] Support vision tower LoRA (#42662)

Signed-off-by: bk-201 <joy25810@foxmail.com> Signed-off-by: linitra24 <Joy25810@foxmail.com> Signed-off-by: linitra24 <joy25810@gmail.com> Co-authored-by: Jee Jee Li <pandaleefree@gmail.com> Co-authored-by: linitra24 <joy25810@gmail.com>

linitra24 committed Aug 13, 2026 at 22:07 UTC 95c9144424060a3ec11090050318a5e702a70be2

10 files changed +309 -48

.buildkite/test_areas/lora.yaml

@@ -48,3 +48,4 @@ steps:

48 - pytest -v -s -x lora/test_olmoe_tp.py

49 - pytest -v -s -x lora/test_gptoss_tp.py

50 - pytest -v -s -x lora/test_qwen35_densemodel_lora.py

51 + - pytest -v -s -x lora/test_gemma4_tp.py

docs/models/supported_models.md

+1 -1

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

523 | `Exaone4_5_ForConditionalGeneration` | EXAONE-4.5 | T + IE+ | `LGAI-EXAONE/EXAONE-4.5-33B`, etc. | ✅︎ | ✅︎ |

524 | `Gemma3ForConditionalGeneration` | Gemma 3 | T + IE+ | `google/gemma-3-4b-it`, `google/gemma-3-27b-it`, etc. | ✅︎ | ✅︎ |

526 +| `Gemma4ForConditionalGeneration` | Gemma 4 | T + I+ + V + A* | `google/gemma-4-E2B-it`, etc. | ✅︎ | ✅︎ |

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

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

tests/lora/conftest.py

@@ -253,6 +253,11 @@ def qwen3vl_vision_lora_files():

253 )

254

255

256 +@pytest.fixture(scope="session")

257 +def gemma4_vision_lora_files():

258 + return hf_api().snapshot_download(repo_id="EpochEcho/gemma4-e2b-it-lora-pokemon")

259 +

260 +

261 @pytest.fixture(scope="session")

262 def qwen3_meowing_lora_files():

263 """Download Qwen3 Meow LoRA files once per test session."""

tests/lora/test_gemma4_tp.py new

+106

@@ -0,0 +1,106 @@

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

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

3 +

4 +# NOTE To avoid overloading the CI pipeline, this test script will not

5 +# be triggered on CI and is primarily intended for local testing and verification.

6 +

7 +import vllm

8 +from vllm.assets.image import ImageAsset

9 +from vllm.lora.request import LoRARequest

10 +

11 +from ..utils import multi_gpu_test

12 +

13 +MODEL_PATH = "google/gemma-4-E2B-it"

14 +

15 +PROMPT_TEMPLATE = """<bos><|turn>user

16 +<|image|>What is in the image?<turn|>

17 +<|turn>model

18 +"""

19 +

20 +TEST_IMAGES = [

21 + ImageAsset("stop_sign"),

22 + ImageAsset("cherry_blossom"),

23 +]

24 +

25 +EXPECTED_OUTPUTS_VISION = [

26 + "A red stop sign stands prominently in the foreground.",

27 + "A majestic skyscraper stands tall, partially obscured by a vibrant "

28 + "canopy of cherry blossoms, against a clear blue sky.",

29 +]

30 +

31 +

32 +def generate_and_test(llm: vllm.LLM, lora_path: str, lora_id: int) -> None:

33 + prompts = [

34 + {

35 + "prompt": PROMPT_TEMPLATE,

36 + "multi_modal_data": {"image": asset.pil_image},

37 + }

38 + for asset in TEST_IMAGES

39 + ]

40 + sampling_params = vllm.SamplingParams(temperature=0, max_tokens=128)

41 + outputs = llm.generate(

42 + prompts,

43 + sampling_params,

44 + lora_request=LoRARequest(str(lora_id), lora_id, lora_path),

45 + )

46 +

47 + generated_texts = [output.outputs[0].text.strip() for output in outputs]

48 + for generated, expected in zip(generated_texts, EXPECTED_OUTPUTS_VISION):

49 + assert generated.startswith(expected), (

50 + f"Generated text {generated!r} does not match expected {expected!r}"

51 + )

52 +

53 +

54 +def test_gemma4_lora(gemma4_vision_lora_files):

55 + llm = vllm.LLM(

56 + MODEL_PATH,

57 + max_model_len=1024,

58 + enable_lora=True,

59 + enforce_eager=True,

60 + max_loras=4,

61 + trust_remote_code=True,

62 + limit_mm_per_prompt={"image": 1},

63 + mm_processor_cache_gb=0,

64 + enable_tower_connector_lora=True,

65 + )

66 +

67 + generate_and_test(llm, gemma4_vision_lora_files, lora_id=1)

68 + generate_and_test(llm, gemma4_vision_lora_files, lora_id=2)

69 +

70 +

71 +@multi_gpu_test(num_gpus=2)

72 +def test_gemma4_lora_tp2(gemma4_vision_lora_files):

73 + llm = vllm.LLM(

74 + MODEL_PATH,

75 + max_model_len=1024,

76 + enforce_eager=True,

77 + enable_lora=True,

78 + max_loras=4,

79 + trust_remote_code=True,

80 + tensor_parallel_size=2,

81 + limit_mm_per_prompt={"image": 1},

82 + mm_processor_cache_gb=0,

83 + enable_tower_connector_lora=True,

84 + )

85 +

86 + generate_and_test(llm, gemma4_vision_lora_files, lora_id=1)

87 + generate_and_test(llm, gemma4_vision_lora_files, lora_id=2)

88 +

89 +

90 +@multi_gpu_test(num_gpus=4)

91 +def test_gemma4_lora_tp4(gemma4_vision_lora_files):

92 + llm = vllm.LLM(

93 + MODEL_PATH,

94 + max_model_len=1024,

95 + enforce_eager=True,

96 + enable_lora=True,

97 + max_loras=4,

98 + trust_remote_code=True,

99 + tensor_parallel_size=4,

100 + limit_mm_per_prompt={"image": 1},

101 + mm_processor_cache_gb=0,

102 + enable_tower_connector_lora=True,

103 + )

104 +

105 + generate_and_test(llm, gemma4_vision_lora_files, lora_id=1)

106 + generate_and_test(llm, gemma4_vision_lora_files, lora_id=2)

tests/v1/worker/test_gpu_model_runner.py

+6 -2

@@ -380,9 +380,13 @@ def test_select_common_block_size_no_valid_option():

380

381 def test_set_active_mm_loras_builds_tower_and_connector_mappings():

382 model = Mock()

383 - model.get_num_mm_encoder_tokens.side_effect = lambda num_embeds: num_embeds + 1

383 + model.get_mm_lora_token_counts.side_effect = (

384 + lambda *, modality, mm_kwargs, num_mm_embeds: (

385 + num_mm_embeds + 1,

386 + num_mm_embeds + 11,

387 + )

388 + )

389 model.get_mm_mapping.return_value = SimpleNamespace(connector=True)

385 - model.get_num_mm_connector_tokens.side_effect = lambda num_tokens: num_tokens + 10

390

391 lora_manager = Mock()

392 lora_manager.supports_tower_connector_lora.return_value = True

vllm/lora/model_manager.py

+20 -6

@@ -247,8 +247,17 @@ class LoRAModelManager:

247

248 mm_budget = MultiModalBudget(vllm_config, mm_registry)

249 limit_per_prompt = max(mm_budget.mm_max_items_per_prompt.values())

250 - num_encoder_tokens = self.model.get_num_mm_encoder_tokens(

251 - mm_budget.get_encoder_budget()

250 + max_lora_tokens = mm_budget.get_encoder_budget()

251 + lora_token_counts_by_modality = [

252 + self.model.get_mm_lora_token_counts(

253 + modality=modality,

254 + mm_kwargs=None,

255 + num_mm_embeds=max_lora_tokens,

256 + )

257 + for modality in mm_budget.mm_max_toks_per_item

258 + ]

259 + num_encoder_tokens = max(

260 + tower_tokens for tower_tokens, _ in lora_token_counts_by_modality

261 )

262

263 # Tower wrappers

@@ -263,10 +272,15 @@ class LoRAModelManager:

272

273 # Use wrapper for connector if present.

274 if self.mm_mapping.connector:

266 - if hasattr(self.model, "get_num_mm_connector_tokens"):

267 - connector_tokens = self.model.get_num_mm_connector_tokens(

268 - num_encoder_tokens

269 - )

275 + connector_tokens = max(

276 + (

277 + connector_tokens

278 + for _, connector_tokens in lora_token_counts_by_modality

279 + if connector_tokens is not None

280 + ),

281 + default=None,

282 + )

283 + if connector_tokens is not None:

284 connector_punica_wrapper = get_punica_wrapper(

285 connector_tokens,

286 max_batches=self.max_num_seqs * limit_per_prompt,

vllm/model_executor/models/gemma4_mm.py

+120 -14

@@ -50,6 +50,7 @@ from vllm.model_executor.models.transformers.utils import recursive_replace_line

50 from vllm.multimodal import MULTIMODAL_REGISTRY

51 from vllm.multimodal.inputs import (

52 MultiModalFieldConfig,

53 + MultiModalKwargsItem,

54 MultiModalKwargsItems,

55 VideoItem,

56 )

@@ -1037,6 +1038,10 @@ class Gemma4ForConditionalGeneration(

1038 self.multimodal_config = multimodal_config

1039 self.model_dtype = vllm_config.model_config.dtype

1040 self.vllm_config = vllm_config

1041 + lora_config = vllm_config.lora_config

1042 + self._enable_mm_lora = bool(

1043 + lora_config is not None and lora_config.enable_tower_connector_lora

1044 + )

1045

1046 # Only quantize towers when the quant method supports their

1047 # dimensions. BNB/torchao handle arbitrary sizes; other methods

@@ -1265,8 +1270,8 @@ class Gemma4ForConditionalGeneration(

1270

1271 Groups images by patch count (resolution bucket) so each

1272 encoder call processes a uniform-shape batch with no

1268 - cross-resolution padding. Pooling and projection are then

1269 - applied over a single concatenated tensor for all images.

1273 + cross-resolution padding. With MM LoRA enabled, all images are

1274 + padded into one batch so the encoder call matches the tower mapping.

1275 """

1276 pixel_values = image_input["pixel_values"]

1277 pixel_position_ids = image_input["pixel_position_ids"]

@@ -1284,24 +1289,65 @@ class Gemma4ForConditionalGeneration(

1289 if isinstance(pixel_values, list)

1290 else pixel_values.shape[0]

1291 )

1292 + pool_position_ids = pixel_position_ids

1293 +

1294 + if self._enable_mm_lora:

1295 + max_soft_tokens = vision_cfg.default_output_length

1296 + mm_processor_kwargs = getattr(

1297 + getattr(self, "multimodal_config", None),

1298 + "mm_processor_kwargs",

1299 + None,

1300 + )

1301 + if isinstance(mm_processor_kwargs, Mapping):

1302 + value, _ = _get_max_soft_tokens(mm_processor_kwargs)

1303 + if isinstance(value, int) and value in _SUPPORTED_SOFT_TOKENS:

1304 + max_soft_tokens = value

1305 +

1306 + max_patches = max_soft_tokens * pooling_k2

1307 + padded_position_ids: list[torch.Tensor] = []

1308 + for idx in range(total_images):

1309 + pv = pixel_values[idx]

1310 + pp = pixel_position_ids[idx]

1311 + num_patches = pv.shape[0]

1312 + if num_patches > max_patches:

1313 + raise ValueError(

1314 + f"Image {idx} has {num_patches} patches, which exceeds "

1315 + f"the MM LoRA patch limit of {max_patches}."

1316 + )

1317

1288 - for idx in range(total_images):

1289 - pv = pixel_values[idx]

1290 - pp = pixel_position_ids[idx]

1291 - buckets.setdefault(pv.shape[0], []).append((idx, pv, pp))

1318 + pad_len = max_patches - num_patches

1319 + pv = torch.cat(

1320 + (pv, pv.new_zeros((pad_len, *pv.shape[1:]))),

1321 + dim=0,

1322 + )

1323 + pp = torch.cat(

1324 + (pp, pp.new_full((pad_len, *pp.shape[1:]), -1)),

1325 + dim=0,

1326 + )

1327 + buckets.setdefault(max_patches, []).append((idx, pv, pp))

1328 + padded_position_ids.append(pp)

1329 + pool_position_ids = padded_position_ids

1330 + else:

1331 + for idx in range(total_images):

1332 + pv = pixel_values[idx]

1333 + pp = pixel_position_ids[idx]

1334 + buckets.setdefault(pv.shape[0], []).append((idx, pv, pp))

1335

1336 # Encode each resolution bucket in memory-safe chunks. Re-read

1337 # free memory per bucket because the previous bucket's encoder

1338 # pass has already allocated activations we should account for.

1339 last_hidden_states_map: dict[int, torch.Tensor] = {}

1340 for patches, items in buckets.items():

1298 - free, total = torch.accelerator.get_memory_info()

1299 - max_batch_size = min(

1300 - len(items),

1301 - self._encoder_chunk(

1302 - patches, free, total, vision_cfg.position_embedding_size

1303 - ),

1304 - )

1341 + if self._enable_mm_lora:

1342 + max_batch_size = len(items)

1343 + else:

1344 + free, total = torch.accelerator.get_memory_info()

1345 + max_batch_size = min(

1346 + len(items),

1347 + self._encoder_chunk(

1348 + patches, free, total, vision_cfg.position_embedding_size

1349 + ),

1350 + )

1351

1352 for chunk_idx in range(0, len(items), max_batch_size):

1353 chunk_items = items[chunk_idx : chunk_idx + max_batch_size]

@@ -1338,7 +1384,7 @@ class Gemma4ForConditionalGeneration(

1384 output_length = chunk_hidden.shape[0] // pooling_k2

1385

1386 single_hidden = chunk_hidden.unsqueeze(0)

1341 - single_pos_ids = pixel_position_ids[orig_idx].unsqueeze(0)

1387 + single_pos_ids = pool_position_ids[orig_idx].unsqueeze(0)

1388 padding_positions = (single_pos_ids == -1).all(dim=-1)

1389

1390 pooled_states, valid_mask = vt.pooler(

@@ -2146,6 +2192,66 @@ class Gemma4ForConditionalGeneration(

2192 tower_model=tower_models,

2193 )

2194

2195 + def get_mm_lora_token_counts(

2196 + self,

2197 + *,

2198 + modality: str,

2199 + mm_kwargs: MultiModalKwargsItem | None,

2200 + num_mm_embeds: int,

2201 + ) -> tuple[int, int | None]:

2202 + if modality in ("image", "video"):

2203 + vision_config = self.config.vision_config

2204 + pooling_k2 = vision_config.pooling_kernel_size**2

2205 +

2206 + if modality == "image":

2207 + pixel_values_key = "pixel_values"

2208 + max_soft_tokens = vision_config.default_output_length

2209 + mm_processor_kwargs = getattr(

2210 + getattr(self, "multimodal_config", None),

2211 + "mm_processor_kwargs",

2212 + None,

2213 + )

2214 + if isinstance(mm_processor_kwargs, Mapping):

2215 + val, _ = _get_max_soft_tokens(mm_processor_kwargs)

2216 + if isinstance(val, int) and val in _SUPPORTED_SOFT_TOKENS:

2217 + max_soft_tokens = val

2218 + else:

2219 + pixel_values_key = "pixel_values_videos"

2220 + max_soft_tokens = _VIDEO_MAX_SOFT_TOKENS

2221 +

2222 + tower_tokens = max_soft_tokens * pooling_k2 if modality == "image" else None

2223 + connector_tokens = num_mm_embeds

2224 + if tower_tokens is None and mm_kwargs is not None:

2225 + field = mm_kwargs.get(pixel_values_key)

2226 + if field is not None:

2227 + data = field.data

2228 + if isinstance(data, torch.Tensor) and data.ndim >= 2:

2229 + tower_tokens = int(math.prod(data.shape[:-1]))

2230 +

2231 + if tower_tokens is None:

2232 + min_soft_tokens = min(_SUPPORTED_SOFT_TOKENS)

2233 + tower_tokens = (

2234 + math.ceil(num_mm_embeds / min_soft_tokens)

2235 + * max_soft_tokens

2236 + * pooling_k2

2237 + )

2238 +

2239 + if modality == "audio":

2240 + tower_tokens = num_mm_embeds

2241 + connector_tokens = num_mm_embeds

2242 +

2243 + if mm_kwargs is not None:

2244 + field = mm_kwargs.get("input_features_padded")

2245 + if field is not None:

2246 + data = field.data

2247 + if isinstance(data, torch.Tensor) and data.ndim >= 2:

2248 + batch_size = math.prod(data.shape[:-2])

2249 + audio_tokens = batch_size * math.ceil(data.shape[-2] / 4)

2250 + tower_tokens = audio_tokens

2251 + connector_tokens = audio_tokens

2252 +

2253 + return tower_tokens, connector_tokens

2254 +

2255 @classmethod

2256 def get_placeholder_str(cls, modality: str, i: int) -> str | None:

2257 if modality == "image":

vllm/model_executor/models/interfaces.py

+23 -1

@@ -48,7 +48,7 @@ if TYPE_CHECKING:

48 from vllm.model_executor.layers.mamba.mamba_utils import MambaStateCopyFunc

49 from vllm.model_executor.models.interfaces_base import VllmModel

50 from vllm.model_executor.models.utils import WeightsMapper

51 - from vllm.multimodal.inputs import MultiModalFeatureSpec

51 + from vllm.multimodal.inputs import MultiModalFeatureSpec, MultiModalKwargsItem

52 from vllm.multimodal.registry import _ProcessorFactories

53 from vllm.sequence import IntermediateTensors

54 from vllm.tasks import ScoreType

@@ -385,6 +385,28 @@ class SupportsMultiModal(SupportsMultiModalEmbeddings, Protocol):

385 """

386 ...

387

388 + def get_mm_lora_token_counts(

389 + self,

390 + *,

391 + modality: str,

392 + mm_kwargs: "MultiModalKwargsItem | None",

393 + num_mm_embeds: int,

394 + ) -> tuple[int, int | None]:

395 + """

396 + Return ``(tower_tokens, connector_tokens)`` for multimodal LoRA mappings.

397 +

398 + MM LoRA uses these counts to build adapter mappings for the tower and

399 + connector forwards. Models with multiple modalities can override this

400 + when each modality has different encoder padding or pooling behavior.

401 + """

402 + del modality, mm_kwargs

403 + num_encoder_tokens = self.get_num_mm_encoder_tokens(num_mm_embeds)

404 + num_connector_tokens = self.get_num_mm_connector_tokens(num_encoder_tokens)

405 + return (

406 + num_encoder_tokens,

407 + num_connector_tokens if isinstance(num_connector_tokens, int) else None,

408 + )

409 +

410 @overload

411 def embed_input_ids(self, input_ids: Tensor) -> Tensor: ...

412

vllm/v1/worker/gpu/mm/lora.py

+13 -11

@@ -29,6 +29,7 @@ def set_active_mm_loras(

29 token_lora_mapping: list[int] = []

30 lora_requests = set()

31 encoder_token_counts: list[int] = []

32 + connector_token_counts: list[int] = []

34 # iterate through images

35 for req_id, encoder_input_ids in scheduled_encoder_inputs.items():

@@ -42,10 +43,17 @@ def set_active_mm_loras(

43 # iterate through visual tokens

44 for mm_input_id in encoder_input_ids:

45 pos_info = mm_features[mm_input_id].mm_position

45 - num_tokens = model.get_num_mm_encoder_tokens(pos_info.get_num_embeds())

46 +

47 + tower_tokens, connector_tokens = model.get_mm_lora_token_counts(

48 + modality=mm_features[mm_input_id].modality,

49 + mm_kwargs=mm_features[mm_input_id].data,

50 + num_mm_embeds=pos_info.get_num_embeds(),

51 + )

52 +

53 prompt_lora_mapping.append(lora_id)

47 - token_lora_mapping.extend([lora_id] * num_tokens)

48 - encoder_token_counts.append(num_tokens)

54 + token_lora_mapping.extend([lora_id] * tower_tokens)

55 + encoder_token_counts.append(tower_tokens)

56 + connector_token_counts.append(connector_tokens)

58 if lora_id > 0:

59 lora_request = lora_state.lora_requests.get(req_id)

@@ -69,19 +77,13 @@ def set_active_mm_loras(

77 if (

78 mm_mapping is None

79 or not mm_mapping.connector

72 - or not hasattr(model, "get_num_mm_connector_tokens")

80 + or not all(count is not None for count in connector_token_counts)

81 ):

82 return

84 connector_token_mapping = np.repeat(

85 np.array(prompt_lora_mapping, dtype=np.int32),

78 - np.array(

79 - [

80 - model.get_num_mm_connector_tokens(num_tokens)

81 - for num_tokens in encoder_token_counts

82 - ],

83 - dtype=np.int32,

84 - ),

86 + np.array(connector_token_counts, dtype=np.int32),

87 )

88 lora_manager.set_active_adapters(

89 lora_requests,

vllm/v1/worker/gpu_model_runner.py

+14 -13

@@ -3123,18 +3123,24 @@ class GPUModelRunner(

3123 token_lora_mapping = []

3124 lora_requests = set()

3125 encoder_token_counts = []

3126 + connector_token_counts = []

3127

3127 - for req_id, pos_info in mm_lora_refs:

3128 + for (req_id, pos_info), (modality, mm_item) in zip(

3129 + mm_lora_refs,

3130 + mm_kwargs,

3131 + ):

3132 req_idx = self.input_batch.req_id_to_index[req_id]

3133 lora_id = int(self.input_batch.request_lora_mapping[req_idx])

3134

3131 - # Prefer pos_info.get_num_embeds to count precise MM embedding tokens.

3132 - num_tokens = self.model.get_num_mm_encoder_tokens( # type: ignore[attr-defined]

3133 - pos_info.get_num_embeds()

3135 + tower_tokens, connector_tokens = self.model.get_mm_lora_token_counts( # type: ignore[attr-defined]

3136 + modality=modality,

3137 + mm_kwargs=mm_item,

3138 + num_mm_embeds=pos_info.get_num_embeds(),

3139 )

3140 prompt_lora_mapping.append(lora_id)

3136 - token_lora_mapping.extend([lora_id] * num_tokens)

3137 - encoder_token_counts.append(num_tokens)

3141 + token_lora_mapping.extend([lora_id] * tower_tokens)

3142 + encoder_token_counts.append(tower_tokens)

3143 + connector_token_counts.append(connector_tokens)

3144

3145 if lora_id > 0:

3146 lora_request = self.input_batch.lora_id_to_lora_request.get(lora_id)

@@ -3162,16 +3168,11 @@ class GPUModelRunner(

3168 if (

3169 mm_mapping is not None

3170 and mm_mapping.connector

3165 - and hasattr(self.model, "get_num_mm_connector_tokens")

3171 + and all(count is not None for count in connector_token_counts)

3172 ):

3167 - post_op_counts = [

3168 - self.model.get_num_mm_connector_tokens(num_tokens) # type: ignore[attr-defined]

3169 - for num_tokens in encoder_token_counts

3170 - ]

3171 -

3173 connector_token_mapping = np.repeat(

3174 np.array(prompt_lora_mapping, dtype=np.int32),

3174 - np.array(post_op_counts, dtype=np.int32),

3175 + np.array(connector_token_counts, dtype=np.int32),

3176 )

3177 connector_mapping = LoRAMapping(

3178 index_mapping=tuple(connector_token_mapping.tolist()),