2b75a73b8e · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -61,6 +61,7 @@ from vllm.multimodal.processing.processor import (

61 PromptUpdate,

62 PromptUpdateDetails,

63 )

64 +from vllm.platforms import current_platform

65 from vllm.sequence import IntermediateTensors

66 from vllm.utils.tensor_schema import TensorSchema, TensorShape

67

@@ -960,6 +961,9 @@ class Gemma4ForConditionalGeneration(

961 self.embed_vision = Gemma4MultimodalEmbedder(

962 config.vision_config, config.text_config

963 )

964 + # Lazy-initialized on first encoder call (see _encoder_max_batch).

965 + self._encoder_budget_bytes = 0

966 + self._encoder_bytes_per_patch = 0

967

968 # ---- Audio tower (variants with audio_config) ----

969 if config.audio_config is not None:

@@ -1100,6 +1104,19 @@ class Gemma4ForConditionalGeneration(

1104 )

1105 return mm_input_by_modality

1106

1107 + def _encoder_max_batch(self, patches_per_item: int) -> int:

1108 + """Max items per encoder call given per-item patch count."""

1109 + if self._encoder_budget_bytes == 0:

1110 + total_mem = current_platform.get_device_total_memory()

1111 + self._encoder_budget_bytes = int(total_mem * 0.05)

1112 + logger.info(

1113 + "Encoder memory budget: %.1fGB (total=%.1fGB)",

1114 + self._encoder_budget_bytes / 1024**3,

1115 + total_mem / 1024**3,

1116 + )

1117 + cost = patches_per_item * self._encoder_bytes_per_patch

1118 + return max(1, self._encoder_budget_bytes // cost) if cost > 0 else 1

1119 +

1120 # ------------------------------------------------------------------ #

1121 # Image processing

1122 # ------------------------------------------------------------------ #

@@ -1108,73 +1125,103 @@ class Gemma4ForConditionalGeneration(

1125 self,

1126 image_input: Gemma4ImageInputs,

1127 ) -> list[torch.Tensor]:

1128 + """Batch-encode images through the vision tower.

1129 +

1130 + Groups images by patch count (resolution bucket) so each

1131 + encoder call processes a uniform-shape batch with no

1132 + cross-resolution padding. Pooling and projection are then

1133 + applied over a single concatenated tensor for all images.

1134 + """

1135 pixel_values = image_input["pixel_values"]

1136 pixel_position_ids = image_input["pixel_position_ids"]

1137

1114 - # The HF image processor now outputs pre-patchified data:

1115 - # pixel_values: (num_images, max_patches, patch_pixels)

1116 - # pixel_position_ids: (num_images, max_patches, 2)

1117 - # We call the vision tower's forward() directly, which handles

1118 - # patch embedding, encoding, pooling, padding removal, and

1119 - # optional standardization internally.

1138 vt = self.vision_tower

1139 pooling_k2 = self.config.vision_config.pooling_kernel_size**2

1140

1123 - # TODO: Move this per-image loop into the input processor to

1124 - # reduce dynamism at the model runner / engine core. This

1125 - # requires spatially padding all images to uniform (H_max,

1126 - # W_max) in _call_hf_processor() so they arrive as a single

1127 - # stacked tensor, tracking padded regions via image_sizes

1128 - # metadata, and validating numerical equivalence with the

1129 - # current per-image path.

1130 - #

1141 # Concurrent requests with different image resolutions may

1142 # arrive as a list of per-image tensors, while same-resolution

1133 - # batches may arrive as a stacked tensor. Both forms are

1134 - # iterable over the per-image dimension.

1135 -

1136 - # Process each image individually through the vision tower.

1137 - # The vision tower's forward() strips padding and returns a

1138 - # flat tensor of valid tokens. We process per-image to get

1139 - # variable-length outputs matching the dynamic token count

1140 - # from get_image_repl.

1141 - per_image_features = []

1142 - for pv, pp in zip(pixel_values, pixel_position_ids, strict=True):

1143 - pv = pv.unsqueeze(0) # (1, max_patches, patch_pixels)

1144 - pp = pp.unsqueeze(0) # (1, max_patches, 2)

1145 -

1146 - # Derive the pooler's output_length from the total patch

1147 - # count (including padding). The vision tower encoder

1148 - # processes ALL patches — padding patches get zero hidden

1149 - # states but still occupy sequence positions. The pooler's

1150 - # _avg_pool_by_positions requires:

1151 - # input_seq_len / output_length == k²

1152 - # where k == pooling_kernel_size. The image processor

1153 - # allocates max_patches = max_soft_tokens * k² total slots,

1154 - # so output_length = max_patches / k² == max_soft_tokens.

1155 - # Without this, the pooler falls back to

1156 - # config.image_seq_length (e.g. 280), which fails when a

1157 - # different max_soft_tokens was used at preprocessing time.

1158 - max_patches = pv.shape[1]

1159 - output_length = max_patches // pooling_k2

1160 -

1161 - vt_output = vt(pv, pp, output_length=output_length)

1162 - # last_hidden_state: (num_valid_tokens, hidden_size)

1163 - # — already flat with padding stripped by the vision tower

1164 - per_image_features.append(vt_output.last_hidden_state)

1165 -

1166 - # Project each image's features into LM embedding space.

1167 - # Per-image loop is required because images have variable

1168 - # token counts after padding removal.

1169 - # Cast to match the projection layer's dtype (model may be

1170 - # bf16 while the vision tower outputs fp32).

1171 - target_dtype = self.embed_vision.embedding_projection.weight.dtype

1172 - return [

1173 - self.embed_vision(inputs_embeds=img.unsqueeze(0).to(target_dtype)).squeeze(

1174 - 0

1143 + # batches may arrive as a stacked tensor.

1144 + buckets: dict[int, list[tuple[int, torch.Tensor, torch.Tensor]]] = {}

1145 + total_images = (

1146 + len(pixel_values)

1147 + if isinstance(pixel_values, list)

1148 + else pixel_values.shape[0]

1149 + )

1150 +

1151 + for idx in range(total_images):

1152 + pv = pixel_values[idx]

1153 + pp = pixel_position_ids[idx]

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

1155 +

1156 + # Encode each resolution bucket in memory-safe chunks.

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

1158 + for patches, items in buckets.items():

1159 + max_batch_size = min(len(items), self._encoder_max_batch(patches))

1160 +

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

1162 + chunk_items = items[chunk_idx : chunk_idx + max_batch_size]

1163 +

1164 + pv_tensor = torch.cat(

1165 + [item[1].unsqueeze(0) for item in chunk_items], dim=0

1166 + )

1167 + pp_tensor = torch.cat(

1168 + [item[2].unsqueeze(0) for item in chunk_items], dim=0

1169 + )

1170 + pad_tensor = (pp_tensor == -1).all(dim=-1)

1171 +

1172 + inputs_embeds = vt.patch_embedder(pv_tensor, pp_tensor, pad_tensor)

1173 + encoder_outputs = vt.encoder(

1174 + inputs_embeds=inputs_embeds,

1175 + attention_mask=~pad_tensor,

1176 + pixel_position_ids=pp_tensor,

1177 + )

1178 + hidden_states = encoder_outputs.last_hidden_state

1179 +

1180 + for i, (orig_idx, _, _) in enumerate(chunk_items):

1181 + last_hidden_states_map[orig_idx] = hidden_states[i]

1182 +

1183 + # Pool per image to strip padding and reduce spatial resolution.

1184 + all_valid_states: list[torch.Tensor] = [None] * total_images # type: ignore[list-item]

1185 + valid_lens = [0] * total_images

1186 +

1187 + for orig_idx in range(total_images):

1188 + chunk_hidden = last_hidden_states_map[orig_idx]

1189 + output_length = chunk_hidden.shape[0] // pooling_k2

1190 +

1191 + single_hidden = chunk_hidden.unsqueeze(0)

1192 + single_pos_ids = pixel_position_ids[orig_idx].unsqueeze(0)

1193 + padding_positions = (single_pos_ids == -1).all(dim=-1)

1194 +

1195 + pooled_states, valid_mask = vt.pooler(

1196 + hidden_states=single_hidden,

1197 + pixel_position_ids=single_pos_ids,

1198 + padding_positions=padding_positions,

1199 + output_length=output_length,

1200 )

1176 - for img in per_image_features

1177 - ]

1201 + valid_states = pooled_states[valid_mask]

1202 +

1203 + if getattr(vt.config, "standardize", False):

1204 + valid_states = (valid_states - vt.std_bias) * vt.std_scale

1205 +

1206 + all_valid_states[orig_idx] = valid_states

1207 + valid_lens[orig_idx] = valid_states.shape[0]

1208 +

1209 + target_dtype = self.embed_vision.embedding_projection.weight.dtype

1210 +

1211 + # Project all images in a single batched call.

1212 + flat_valid_states = torch.cat(all_valid_states, dim=0).to(target_dtype)

1213 + flat_proj_embs = self.embed_vision(

1214 + inputs_embeds=flat_valid_states.unsqueeze(0)

1215 + ).squeeze(0)

1216 +

1217 + # Split back into per-image tensors (slicing returns views).

1218 + per_image_embeddings: list[torch.Tensor] = []

1219 + offset = 0

1220 + for length in valid_lens:

1221 + per_image_embeddings.append(flat_proj_embs[offset : offset + length])

1222 + offset += length

1223 +

1224 + return per_image_embeddings

1225

1226 # ------------------------------------------------------------------ #

1227 # Video processing (frames through vision tower)

@@ -1184,16 +1231,16 @@ class Gemma4ForConditionalGeneration(

1231 self,

1232 video_input: dict[str, torch.Tensor],

1233 ) -> list[torch.Tensor]:

1187 - """Process video frames through the vision tower.

1234 + """Batch-encode video frames through the vision tower.

1235

1189 - Reuses the image processing pipeline — Gemma4 has no separate

1190 - video tower; video frames are just images at lower resolution

1191 - (max_soft_tokens=70).

1236 + Gemma4 has no separate video tower; video frames are images at

1237 + lower resolution (max_soft_tokens=70). All frames across all

1238 + videos in the batch are encoded together in chunks, then pooled

1239 + and projected in a single batched call.

1240

1241 Returns one concatenated embedding tensor per video (not per

1194 - frame), because vLLM treats one video as one multimodal item.

1195 - The flat_from_sizes field config groups all frames of a video

1196 - together, so embed_multimodal must return one tensor per video.

1242 + frame), matching the flat_from_sizes grouping that vLLM expects

1243 + for embed_multimodal.

1244 """

1245 pixel_values = video_input["pixel_values_videos"]

1246 pixel_position_ids = video_input["pixel_position_ids_videos"]

@@ -1203,35 +1250,74 @@ class Gemma4ForConditionalGeneration(

1250 pooling_k2 = self.config.vision_config.pooling_kernel_size**2

1251 target_dtype = self.embed_vision.embedding_projection.weight.dtype

1252

1206 - # Split flat tensors into per-video chunks

1253 if isinstance(frame_counts, torch.Tensor):

1254 fc_list = frame_counts.tolist()

1255 else:

1256 fc_list = list(frame_counts)

1257

1212 - pv_per_video = torch.split(pixel_values, fc_list, dim=0)

1213 - pp_per_video = torch.split(pixel_position_ids, fc_list, dim=0)

1258 + total_frames = pixel_values.shape[0]

1259 + max_batch_size = min(

1260 + total_frames, self._encoder_max_batch(pixel_values.shape[1])

1261 + )

1262

1215 - per_video_embeddings = []

1216 - for pv_chunk, pp_chunk in zip(pv_per_video, pp_per_video):

1217 - frame_embs = []

1218 - for i in range(pv_chunk.shape[0]):

1219 - pv = pv_chunk[i].unsqueeze(0)

1220 - pp = pp_chunk[i].unsqueeze(0)

1263 + padding_positions = (pixel_position_ids == -1).all(dim=-1)

1264

1222 - max_patches = pv.shape[1]

1223 - output_length = max_patches // pooling_k2

1265 + # Encode frames in chunks bounded by _encoder_max_batch.

1266 + last_hidden_states_list: list[torch.Tensor] = []

1267 + for i in range(0, total_frames, max_batch_size):

1268 + pv_chunk = pixel_values[i : i + max_batch_size]

1269 + pp_chunk = pixel_position_ids[i : i + max_batch_size]

1270 + pad_chunk = padding_positions[i : i + max_batch_size]

1271

1225 - vt_output = vt(pv, pp, output_length=output_length)

1226 - frame_emb = self.embed_vision(

1227 - inputs_embeds=(

1228 - vt_output.last_hidden_state.unsqueeze(0).to(target_dtype)

1229 - )

1230 - ).squeeze(0)

1231 - frame_embs.append(frame_emb)

1272 + inputs_embeds = vt.patch_embedder(pv_chunk, pp_chunk, pad_chunk)

1273 + encoder_outputs = vt.encoder(

1274 + inputs_embeds=inputs_embeds,

1275 + attention_mask=~pad_chunk,

1276 + pixel_position_ids=pp_chunk,

1277 + )

1278 + last_hidden_states_list.append(encoder_outputs.last_hidden_state)

1279 +

1280 + last_hidden_states = torch.cat(last_hidden_states_list, dim=0)

1281

1233 - # Concatenate all frames of this video into one tensor.

1234 - per_video_embeddings.append(torch.cat(frame_embs, dim=0))

1282 + # Pool per frame to strip padding and reduce spatial resolution.

1283 + output_length = pixel_values.shape[1] // pooling_k2

1284 + all_frame_valid_states: list[torch.Tensor] = []

1285 + frame_valid_lens: list[int] = []

1286 +

1287 + for i in range(total_frames):

1288 + single_hidden = last_hidden_states[i].unsqueeze(0)

1289 + single_pos_ids = pixel_position_ids[i].unsqueeze(0)

1290 + single_pad_pos = padding_positions[i].unsqueeze(0)

1291 +

1292 + pooled_states, valid_mask = vt.pooler(

1293 + hidden_states=single_hidden,

1294 + pixel_position_ids=single_pos_ids,

1295 + padding_positions=single_pad_pos,

1296 + output_length=output_length,

1297 + )

1298 + valid_states = pooled_states[valid_mask]

1299 +

1300 + if getattr(vt.config, "standardize", False):

1301 + valid_states = (valid_states - vt.std_bias) * vt.std_scale

1302 +

1303 + all_frame_valid_states.append(valid_states)

1304 + frame_valid_lens.append(valid_states.shape[0])

1305 +

1306 + # Project all frames in a single batched call.

1307 + flat_valid_states = torch.cat(all_frame_valid_states, dim=0).to(target_dtype)

1308 + flat_proj_embs = self.embed_vision(

1309 + inputs_embeds=flat_valid_states.unsqueeze(0)

1310 + ).squeeze(0)

1311 +

1312 + # Regroup into per-video tensors (slicing returns views).

1313 + per_video_embeddings: list[torch.Tensor] = []

1314 + frame_idx = 0

1315 + offset = 0

1316 + for count in fc_list:

1317 + video_tokens = sum(frame_valid_lens[frame_idx : frame_idx + count])

1318 + per_video_embeddings.append(flat_proj_embs[offset : offset + video_tokens])

1319 + offset += video_tokens

1320 + frame_idx += count

1321

1322 return per_video_embeddings

1323

@@ -1452,7 +1538,15 @@ class Gemma4ForConditionalGeneration(

1538 self,

1539 ignore_unexpected_prefixes=ignore_prefixes,

1540 )

1455 - return loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1541 + loaded = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1542 +

1543 + # Per-patch activation cost for dynamic encoder batch sizing.

1544 + vis_cfg = self.config.vision_config

1545 + self._encoder_bytes_per_patch = (

1546 + vis_cfg.hidden_size * 2 * vis_cfg.num_hidden_layers

1547 + )

1548 +

1549 + return loaded

1550

1551 # ------------------------------------------------------------------ #

1552 # LoRA / multimodal mapping

[Perf][Gemma4] Batch vision encoder calls for image and video processing (#43169)