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

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 6, 10:45 UTC

Support temporal compression for Nemotron-3-VL videos (#36808)

Signed-off-by: Collin McCarthy <cmccarthy@nvidia.com>

Collin McCarthy committed Mar 19, 2026 at 01:02 UTC 0b6d52629fe84f0071ce41a954162e59dc98157e

5 files changed +553 -130

vllm/model_executor/models/nano_nemotron_vl.py

+175 -35

@@ -8,6 +8,7 @@

8 # --------------------------------------------------------

10 import copy

11 +import math

12 import warnings

13 from abc import abstractmethod

14 from collections.abc import Iterable, Mapping, Sequence

@@ -77,6 +78,7 @@ from vllm.renderers import TokenizeParams

78 from vllm.sequence import IntermediateTensors

79 from vllm.tokenizers import cached_tokenizer_from_config

80 from vllm.transformers_utils.configs.radio import RadioConfig

81 +from vllm.transformers_utils.processors.internvl import get_internvl_target_ratios

82 from vllm.transformers_utils.processors.nano_nemotron_vl import (

83 AUDIO_CONTEXT,

84 IMG_CONTEXT,

@@ -85,7 +87,7 @@ from vllm.transformers_utils.processors.nano_nemotron_vl import (

87 BaseNanoNemotronVLProcessor,

88 DynamicResolutionImageTiler,

89 NanoNemotronVLProcessor,

88 - get_internvl_target_ratios,

90 + get_video_target_size_and_feature_size,

91 )

92 from vllm.utils.tensor_schema import TensorSchema, TensorShape

@@ -295,10 +297,13 @@ class NanoNemotronVLProcessingInfo(BaseNanoNemotronVLProcessingInfo):

297 max_videos = mm_counts.get("video", 0)

298

299 processor = self.get_hf_processor() # we get the CustomProcessor here

300 + T = processor.video_temporal_patch_size

301

302 max_image_tokens = self.get_max_image_tokens() * max_images

300 - max_total_frames = (seq_len - max_image_tokens) // processor.num_image_token

301 - max_frames_per_video = max_total_frames // max(max_videos, 1)

303 + tokens_per_tubelet = processor.num_video_token

304 + max_total_tubelets = (seq_len - max_image_tokens) // tokens_per_tubelet

305 + max_tubelets_per_video = max_total_tubelets // max(max_videos, 1)

306 + max_frames_per_video = max_tubelets_per_video * T

307 return max(max_frames_per_video, 1)

308

309 def get_hf_processor(self, **kwargs: object) -> NanoNemotronVLProcessor:

@@ -589,28 +594,49 @@ class NanoNemotronVLMultiModalProcessor(

594 video_num_patches = []

595

596 def get_video_replacement_internvl(item_idx: int):

592 - feature_size = hf_processor.num_image_token

597 video, metadata = mm_items["video"][item_idx]

598 + patch_size = hf_processor.config.patch_size

599 + downsample_ratio = hf_processor.config.downsample_ratio

600 + target_patches = hf_processor.video_target_num_patches

601 +

602 + if target_patches is not None and video is not None and video.shape[0] > 0:

603 + orig_h, orig_w = video.shape[1], video.shape[2]

604 + _, _, feature_size = get_video_target_size_and_feature_size(

605 + orig_w=orig_w,

606 + orig_h=orig_h,

607 + target_patches=target_patches,

608 + maintain_aspect_ratio=hf_processor.video_maintain_aspect_ratio,

609 + patch_size=patch_size,

610 + downsample_ratio=downsample_ratio,

611 + )

612 + else:

613 + feature_size = hf_processor.num_image_token

614 num_patches = video_num_patches[item_idx]

615 if num_patches is not None:

616 assert isinstance(num_patches, int)

617

618 + T = hf_processor.video_temporal_patch_size

619 + if T > 1 and num_patches is not None:

620 + num_tubelets = math.ceil(num_patches / T)

621 + else:

622 + num_tubelets = num_patches

623 +

624 video_pruning_rate = self.info.ctx.get_mm_config().video_pruning_rate

625 if video_pruning_rate is not None and video_pruning_rate > 0.0:

626 # Start of EVS-specific code

627 num_tokens = compute_retained_tokens_count(

628 tokens_per_frame=feature_size,

603 - num_frames=num_patches,

629 + num_frames=num_tubelets,

630 q=video_pruning_rate,

631 )

632 # Here we just need placeholders that won't actually be replaced -

633 # we just need to make sure the total number of tokens is correct

634 # assign all tokens to the first frame

609 - tokens_per_frame = [num_tokens] + [0] * (num_patches - 1)

635 + tokens_per_frame = [num_tokens] + [0] * (num_tubelets - 1)

636

637 # End of EVS-specific code

638 else:

613 - tokens_per_frame = [feature_size] * num_patches

639 + tokens_per_frame = [feature_size] * num_tubelets

640

641 frame_duration_ms = int(1000 / metadata["fps"])

642 return hf_processor.get_video_repl(

@@ -621,6 +647,7 @@ class NanoNemotronVLMultiModalProcessor(

647 img_start_token_ids=hf_processor._img_start_token_ids,

648 img_end_token_ids=hf_processor._img_end_token_ids,

649 img_context_token_ids=hf_processor._img_context_token_ids,

650 + video_temporal_patch_size=T,

651 )

652

653 if self.info.supports_video:

@@ -745,15 +772,39 @@ class NanoNemotronVLDummyInputsBuilder(

772 if self.info.supports_video:

773 config = self.info.get_hf_config()

774 image_size: int = config.force_image_size

775 + processor = self.info.get_hf_processor()

776 +

777 + # When video_target_num_patches is set the per-frame pixel

778 + # resolution can exceed image_size. Use the actual target

779 + # dimensions so that profiling sees the correct upper bound.

780 + if processor.video_target_num_patches is not None:

781 + target_w, target_h, _ = get_video_target_size_and_feature_size(

782 + orig_w=image_size,

783 + orig_h=image_size,

784 + target_patches=processor.video_target_num_patches,

785 + maintain_aspect_ratio=processor.video_maintain_aspect_ratio,

786 + patch_size=config.patch_size,

787 + downsample_ratio=config.downsample_ratio,

788 + )

789 + video_width, video_height = target_w, target_h

790 + else:

791 + video_width, video_height = image_size, image_size

792 +

793 target_num_frames = self.info.get_num_frames_with_most_features(

794 seq_len, mm_counts

795 )

796 + mm_config = self.info.ctx.get_mm_config()

797 + if num_frames := mm_config.media_io_kwargs.get("video", {}).get(

798 + "num_frames"

799 + ):

800 + assert num_frames > 0

801 + target_num_frames = num_frames

802 num_videos = mm_counts.get("video", 0)

803 video_overrides = mm_options.get("video")

804 dummy_video = {

805 "video": self._get_dummy_videos(

755 - width=image_size,

756 - height=image_size,

806 + width=video_width,

807 + height=video_height,

808 num_frames=target_num_frames,

809 num_videos=num_videos,

810 overrides=video_overrides,

@@ -790,6 +841,9 @@ class NanoNemotronVLDummyInputsBuilder(

841 class NemotronH_Nano_VL_V2(

842 nn.Module, HasInnerState, IsHybrid, SupportsMultiModal, SupportsMultiModalPruning

843 ):

844 + requires_sequential_video_encoding = True

845 + """Temporarily needed for dynamic res video w/ conv3d, doesn't support bs>1 yet"""

846 +

847 @classmethod

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

849 if modality.startswith("image"):

@@ -817,6 +871,11 @@ class NemotronH_Nano_VL_V2(

871 self.image_tag_type = config.image_tag_type

872 self.video_pruning_rate = multimodal_config.video_pruning_rate

873

874 + vision_config = getattr(config, "vision_config", config)

875 + self.video_temporal_patch_size: int = getattr(

876 + vision_config, "video_temporal_patch_size", 1

877 + )

878 +

879 with self._mark_language_model(vllm_config):

880 self.language_model = init_vllm_registered_model(

881 vllm_config=vllm_config,

@@ -838,11 +897,12 @@ class NemotronH_Nano_VL_V2(

897

898 mlp1 = nn.Sequential(

899 RMSNorm(

841 - hidden_size=vit_hidden_size * int(1 / self.downsample_ratio) ** 2,

900 + hidden_size=vit_hidden_size

901 + * int(round(1 / self.downsample_ratio)) ** 2,

902 eps=1e-5,

903 ),

904 nn.Linear(

845 - vit_hidden_size * int(1 / self.downsample_ratio) ** 2,

905 + vit_hidden_size * int(round(1 / self.downsample_ratio)) ** 2,

906 vision_projection_hidden_size,

907 bias=False,

908 ),

@@ -958,19 +1018,37 @@ class NemotronH_Nano_VL_V2(

1018 vit_embeds = self.mlp1(vit_embeds)

1019 return vit_embeds

1020

961 - def extract_feature(self, pixel_values: torch.Tensor):

1021 + def extract_feature(

1022 + self,

1023 + pixel_values: torch.Tensor,

1024 + num_frames: int | None = None,

1025 + ) -> torch.Tensor:

1026 # Process images in a micro-batch of at most 128 frames per call

963 - # This is done on purpose to ensure peak GPU ram usage of huge batch

964 - # (namely for really long videos with EVS ON) won't cause any problems

965 - # as we don't support chunked prefill for video media

966 - micro_batch_size = 128

967 - n = pixel_values.shape[0]

1027 + # This is done on purpose to ensure peak GPU ram usage of huge batch

1028 + # (namely for really long videos with EVS ON) won't cause any problems

1029 + # as we don't support chunked prefill for video media

1030 + # When num_frames is provided and temporal_patch_size > 1, consecutive

1031 + # frames are grouped into tubelets — the batch size must be a multiple

1032 + # of T so chunk boundaries don't split a tubelet.

1033 + N, _C, H, W = pixel_values.shape

1034 +

1035 + T = self.video_temporal_patch_size if num_frames is not None else 1

1036 + micro_batch_size = 128 - (128 % T)

1037 + patch_size = self.patch_size

1038 + H_patches = H // patch_size

1039 + W_patches = W // patch_size

1040 +

1041 vit_embeds_list = []

969 - for i in range(0, n, micro_batch_size):

970 - _, vit_embeds = self.vision_model(pixel_values[i : i + micro_batch_size])

1042 + for i in range(0, N, micro_batch_size):

1043 + chunk = pixel_values[i : i + micro_batch_size]

1044 + if num_frames is not None and T > 1:

1045 + _, vit_embeds = self.vision_model(chunk, num_frames=chunk.shape[0])

1046 + else:

1047 + _, vit_embeds = self.vision_model(chunk)

1048 vit_embeds = vit_embeds.to(dtype=torch.bfloat16)

972 - h = w = int(vit_embeds.shape[1] ** 0.5)

973 - vit_embeds = vit_embeds.reshape(vit_embeds.shape[0], h, w, -1)

1049 + vit_embeds = vit_embeds.reshape(

1050 + vit_embeds.shape[0], H_patches, W_patches, -1

1051 + )

1052 vit_embeds = self.pixel_shuffle(

1053 vit_embeds, scale_factor=self.downsample_ratio

1054 )

@@ -1042,16 +1120,21 @@ class NemotronH_Nano_VL_V2(

1120 ) -> tuple[torch.Tensor, ...]:

1121 """Process video input and create final embeddings with video content

1122 and indicator tokens."""

1045 - # Get video embeddings using the same processing as images

1046 - video_embeddings = self._process_image_input(video_input)

1123 + T = self.video_temporal_patch_size

1124 +

1125 + if T > 1:

1126 + video_embeddings = self._extract_video_embeddings_temporal(video_input)

1127 + else:

1128 + video_embeddings = self._process_image_input(video_input)

1129

1130 final_video_embeddings: tuple[torch.Tensor, ...] = ()

1131

1050 - image_rows = image_cols = self.config.force_image_size

1132 downsample_ratio = self.config.downsample_ratio

1133 patch_size = self.config.patch_size

1053 - rows = int(image_rows * downsample_ratio // patch_size)

1054 - cols = int(image_cols * downsample_ratio // patch_size)

1134 + pixel_values = video_input["pixel_values_flat"]

1135 + frame_h, frame_w = pixel_values.shape[-2], pixel_values.shape[-1]

1136 + rows = int(frame_h * downsample_ratio // patch_size)

1137 + cols = int(frame_w * downsample_ratio // patch_size)

1138 video_pruning_rate = self.video_pruning_rate

1139 video_num_frames = video_input["num_patches"].tolist()

1140 video_frames_indices = video_input["frames_indices"].split(video_num_frames)

@@ -1062,13 +1145,14 @@ class NemotronH_Nano_VL_V2(

1145 num_frames = video_num_frames[i]

1146 frames_indices = video_frames_indices[i].tolist()

1147 frame_duration_ms = video_input["frame_duration_ms"][i].item()

1065 - assert single_video_embeddings.shape[0] % num_frames == 0

1148 + num_tubelets = math.ceil(num_frames / T) if T > 1 else num_frames

1149 + assert single_video_embeddings.shape[0] % num_tubelets == 0

1150

1151 if video_pruning_rate is not None and video_pruning_rate > 0.0:

1152 # Start of EVS-specific code

1153 retention_mask = compute_retention_mask(

1154 single_video_embeddings,

1071 - video_size_thw=(num_frames, rows, cols),

1155 + video_size_thw=(num_tubelets, rows, cols),

1156 spatial_merge_size=1,

1157 q=video_pruning_rate,

1158 )

@@ -1077,14 +1161,14 @@ class NemotronH_Nano_VL_V2(

1161 single_video_embeddings = single_video_embeddings[retention_mask]

1162

1163 # calculate the actual number of retained tokens per frame

1080 - retention_mask_thw = retention_mask.reshape(num_frames, rows, cols)

1164 + retention_mask_thw = retention_mask.reshape(num_tubelets, rows, cols)

1165 num_tokens_per_frame = (

1166 retention_mask_thw.sum(dim=(1, 2)).long().tolist()

1167 )

1168 # End of EVS-specific code

1169 else:

1086 - feature_size = single_video_embeddings.shape[0] // num_frames

1087 - num_tokens_per_frame = [feature_size] * num_frames

1170 + feature_size = single_video_embeddings.shape[0] // num_tubelets

1171 + num_tokens_per_frame = [feature_size] * num_tubelets

1172

1173 final_video_embeddings += (

1174 self._create_final_video_embeddings(

@@ -1092,11 +1176,36 @@ class NemotronH_Nano_VL_V2(

1176 num_tokens_per_frame,

1177 frames_indices,

1178 frame_duration_ms,

1179 + video_temporal_patch_size=T,

1180 ),

1181 )

1182

1183 return final_video_embeddings

1184

1185 + def _extract_video_embeddings_temporal(

1186 + self, video_input: NanoNemotronVLVideoPixelInputs

1187 + ) -> tuple[torch.Tensor, ...]:

1188 + """Extract per-video embeddings with temporal compression.

1189 +

1190 + Each video is processed separately through extract_feature with

1191 + num_frames, which uses the fixed-resolution temporal path in RADIO

1192 + (no attention mask, flash attention).

1193 + """

1194 + pixel_values = video_input["pixel_values_flat"]

1195 + num_frames_per_video = video_input["num_patches"].tolist()

1196 + hidden_size = self.config.text_config.hidden_size

1197 +

1198 + results: list[torch.Tensor] = []

1199 + frame_offset = 0

1200 + for nf in num_frames_per_video:

1201 + video_frames = pixel_values[frame_offset : frame_offset + nf]

1202 + frame_offset += nf

1203 +

1204 + vit_embeds = self.extract_feature(video_frames, num_frames=nf)

1205 + results.append(vit_embeds.view(-1, hidden_size))

1206 +

1207 + return tuple(results)

1208 +

1209 def _process_audio_input(

1210 self, audio_input: NanoNemotronVLAudioFeatureInputs

1211 ) -> tuple[torch.Tensor, ...]:

@@ -1134,6 +1243,7 @@ class NemotronH_Nano_VL_V2(

1243 num_tokens_per_frame: list[int],

1244 frames_indices: list[int],

1245 frame_duration_ms: int,

1246 + video_temporal_patch_size: int = 1,

1247 ) -> torch.Tensor:

1248 """Create final embeddings that combine video embeddings with

1249 text embeddings of indicator tokens.

@@ -1161,6 +1271,7 @@ class NemotronH_Nano_VL_V2(

1271 img_start_token_ids=self._img_start_token_ids,

1272 img_end_token_ids=self._img_end_token_ids,

1273 img_context_token_ids=self._img_context_token_ids,

1274 + video_temporal_patch_size=video_temporal_patch_size,

1275 )

1276

1277 # video_repl.full is a list of token IDs

@@ -1207,8 +1318,27 @@ class NemotronH_Nano_VL_V2(

1318 else:

1319 frames_indices = torch.cat([f.flatten() for f in frames_indices], dim=0)

1320

1210 - frame_duration_ms = frame_duration_ms.flatten()

1211 - expected_h = expected_w = self.config.force_image_size

1321 + if torch.is_tensor(frame_duration_ms):

1322 + frame_duration_ms = frame_duration_ms.flatten()

1323 + else:

1324 + frame_duration_ms = torch.cat(

1325 + [f.flatten() for f in frame_duration_ms], dim=0

1326 + )

1327 +

1328 + if (

1329 + torch.is_tensor(pixel_values_flat_video)

1330 + and pixel_values_flat_video.ndim == 5

1331 + ):

1332 + # batched._reduce_data stacked same-shape videos into

1333 + # [num_videos, nf, 3, H, W]; unstack back to a list so the

1334 + # same-H,W cat path below handles it uniformly.

1335 + pixel_values_flat_video = list(pixel_values_flat_video)

1336 +

1337 + if not torch.is_tensor(pixel_values_flat_video):

1338 + pixel_values_flat_video = torch.cat(pixel_values_flat_video, dim=0)

1339 +

1340 + expected_h = pixel_values_flat_video.shape[-2]

1341 + expected_w = pixel_values_flat_video.shape[-1]

1342 num_frames = video_num_patches[0].item()

1343 resolve_bindings = {"h": expected_h, "w": expected_w, "f": num_frames}

1344

@@ -1361,8 +1491,7 @@ class NemotronH_Nano_VL_V2(

1491

1492 self.language_model.load_weights(llm_weights)

1493 self.vision_model.load_weights(vision_weights)

1364 - if self.sound_encoder is not None:

1365 - assert len(sound_weights) > 0

1494 + if self.sound_encoder is not None and len(sound_weights) > 0:

1495 self.sound_encoder.load_weights(sound_weights)

1496

1497 def get_vit_model_from_radio_config(self, hf_config):

@@ -1375,12 +1504,23 @@ class NemotronH_Nano_VL_V2(

1504 image_size = preferred_resolution[0] if preferred_resolution else 224

1505 patch_size = getattr(hf_config_vision, "patch_size", 16)

1506

1507 + # video_temporal_patch_size and separate_video_embedder are

1508 + # top-level vision_config attributes, not inside args.

1509 + video_temporal_patch_size = getattr(

1510 + hf_config_vision, "video_temporal_patch_size", 1

1511 + )

1512 + separate_video_embedder = getattr(

1513 + hf_config_vision, "separate_video_embedder", True

1514 + )

1515 +

1516 radio_config = RadioConfig(

1517 model_name=model_name,

1518 image_size=image_size,

1519 patch_size=patch_size,

1520 norm_mean=hf_config.norm_mean,

1521 norm_std=hf_config.norm_std,

1522 + video_temporal_patch_size=video_temporal_patch_size,

1523 + separate_video_embedder=separate_video_embedder,

1524 **hf_config_vision.args,

1525 )

1526

vllm/model_executor/models/radio.py

+148 -65

@@ -123,6 +123,8 @@ class ViTPatchGenerator(nn.Module):

123 register_multiple: int | None = None,

124 num_registers: int | None = None,

125 patch_bias: bool = False,

126 + temporal_patch_size: int = 1,

127 + separate_video_embedder: bool = True,

128 device=None,

129 dtype=None,

130 ):

@@ -148,6 +150,7 @@ class ViTPatchGenerator(nn.Module):

150 self.patch_size = patch_size

151 self.abs_pos = abs_pos

152 self.embed_dim = embed_dim

153 + self.temporal_patch_size = temporal_patch_size

154

155 self.num_rows = max_input_dims[0] // patch_size

156 self.num_cols = max_input_dims[1] // patch_size

@@ -160,6 +163,21 @@ class ViTPatchGenerator(nn.Module):

163 patch_size, embed_dim, bias=patch_bias, **factory

164 )

165

166 + if temporal_patch_size > 1:

167 + if not separate_video_embedder:

168 + raise NotImplementedError(

169 + "Only separate_video_embedder=True is supported for"

170 + " temporal compression (temporal_patch_size > 1)"

171 + )

172 + self.video_embedder = ViTPatchLinear(

173 + patch_size,

174 + embed_dim,

175 + bias=patch_bias,

176 + temporal_patch_size=temporal_patch_size,

177 + **factory,

178 + )

179 + self._video_embedder_loaded = False

180 +

181 if abs_pos:

182 scale = embed_dim**-0.5

183 self.pos_embed = nn.Parameter(

@@ -196,6 +214,60 @@ class ViTPatchGenerator(nn.Module):

214 return patches, pos_enc

215 return patches

216

217 + def forward_video(self, x: torch.Tensor) -> torch.Tensor:

218 + """Process video frames with temporal compression.

219 +

220 + Groups T consecutive frames into tubelets before embedding.

221 +

222 + Args:

223 + x: [num_frames, 3, H, W] tensor of video frames

224 +

225 + Returns:

226 + Embedded patches with temporal compression applied.

227 + """

228 + if not self._video_embedder_loaded:

229 + raise ValueError(

230 + "Temporal compression (video_temporal_patch_size > 1) requires "

231 + "video_embedder weights, but they were never loaded. "

232 + "Ensure the checkpoint was trained with temporal compression."

233 + )

234 + T = self.temporal_patch_size

235 + input_size = x.shape[2:]

236 +

237 + patches = self.im_to_patches(x) # [N, num_patches, 3*P*P]

238 + num_frames, num_spatial, feat_dim = patches.shape

239 +

240 + # Pad to a multiple of T by repeating the last frame so that

241 + # all tubelets have exactly T frames.

242 + num_pad_frames = (-num_frames) % T

243 + if num_pad_frames > 0:

244 + last_frame_dup = patches[-1:].expand(num_pad_frames, -1, -1)

245 + patches = torch.cat([patches, last_frame_dup], dim=0)

246 +

247 + # Group T frames per tubelet: for each spatial position, concatenate

248 + # features across T consecutive frames; order follows Megatron training

249 + num_frames_padded = patches.shape[0]

250 + num_tublets = num_frames_padded // T

251 + patches = rearrange(

252 + patches,

253 + "(tubelets frames) spatial feat -> tubelets spatial (frames feat)",

254 + tubelets=num_tublets,

255 + frames=T,

256 + spatial=num_spatial,

257 + feat=feat_dim,

258 + )

259 +

260 + patches = self.video_embedder(patches)

261 +

262 + patches, pos_enc = self.apply_pos_enc(patches, input_size=input_size)

263 +

264 + patches = self.cls_token(patches)

265 +

266 + patches = self.patch_normalizer(patches)

267 + if self.return_pos_enc:

268 + return patches, pos_enc

269 + return patches

270 +

271 def apply_pos_enc_dynamic(

272 self, patches: torch.Tensor, imgs_sizes: list[tuple[int, int]]

273 ) -> tuple[torch.Tensor, torch.Tensor | None]:

@@ -381,66 +453,21 @@ class ViTPatchGenerator(nn.Module):

453 return pos_embed

454

455 if self.cpe_mode:

384 - if self.training:

385 - min_scale = math.sqrt(0.1)

386 - scale = (

387 - torch.rand(batch_size, 1, 1, device=pos_embed.device)

388 - * (1 - min_scale)

389 - + min_scale

390 - )

391 - aspect_min = math.log(3 / 4)

392 - aspect_max = -aspect_min

393 - aspect = torch.exp(

394 - torch.rand(batch_size, 1, 1, device=pos_embed.device)

395 - * (aspect_max - aspect_min)

396 - + aspect_min

397 - )

398 -

399 - scale_x = scale * aspect

400 - scale_y = scale * (1 / aspect)

401 - scale_xy = torch.stack([scale_x, scale_y], dim=-1).clamp_(0, 1)

402 -

403 - pos_xy = torch.rand(batch_size, 1, 1, 2, device=pos_embed.device) * (

404 - 1 - scale_xy

405 - )

456 + max_dim = max(input_dims)

457 + pos_embed = F.interpolate(

458 + pos_embed.float(),

459 + size=(max_dim, max_dim),

460 + align_corners=False,

461 + mode="bilinear",

462 + ).to(pos_embed.dtype)

463

407 - lin_x = torch.linspace(

408 - 0, 1, steps=input_dims[1], device=pos_embed.device

409 - )[None, None].expand(batch_size, input_dims[0], -1)

410 - lin_y = torch.linspace(

411 - 0, 1, steps=input_dims[0], device=pos_embed.device

412 - )[None, :, None].expand(batch_size, -1, input_dims[1])

413 -

414 - lin_xy = torch.stack([lin_x, lin_y], dim=-1)

415 -

416 - grid_xy = lin_xy * scale_xy + pos_xy

417 -

418 - # Convert to [-1, 1] range

419 - grid_xy.mul_(2).sub_(1)

420 -

421 - pos_embed = F.grid_sample(

422 - pos_embed.float().expand(batch_size, -1, -1, -1),

423 - grid=grid_xy,

424 - mode="bilinear",

425 - padding_mode="zeros",

426 - align_corners=True,

427 - ).to(pos_embed.dtype)

428 - else:

429 - max_dim = max(input_dims)

430 - pos_embed = F.interpolate(

431 - pos_embed.float(),

432 - size=(max_dim, max_dim),

433 - align_corners=True,

434 - mode="bilinear",

435 - ).to(pos_embed.dtype)

436 -

437 - pos_embed = window_select(pos_embed)

464 + pos_embed = window_select(pos_embed)

465 else:

466 pos_embed = window_select(pos_embed)

467

468 if pos_embed.shape[-2:] != input_dims:

469 pos_embed = F.interpolate(

443 - pos_embed.float(), size=input_dims, align_corners=True, mode="bilinear"

470 + pos_embed.float(), size=input_dims, align_corners=False, mode="bilinear"

471 ).to(pos_embed.dtype)

472

473 pos_embed = pos_embed.flatten(2).permute(0, 2, 1)

@@ -473,9 +500,19 @@ class Im2Patches(nn.Module):

500

501

502 class ViTPatchLinear(nn.Linear):

476 - def __init__(self, patch_size: int, embed_dim: int, bias: bool = False, **factory):

477 - super().__init__(3 * (patch_size**2), embed_dim, bias=bias, **factory)

503 + def __init__(

504 + self,

505 + patch_size: int,

506 + embed_dim: int,

507 + bias: bool = False,

508 + temporal_patch_size: int = 1,

509 + **factory,

510 + ):

511 + super().__init__(

512 + 3 * temporal_patch_size * (patch_size**2), embed_dim, bias=bias, **factory

513 + )

514 self.patch_size = patch_size

515 + self.temporal_patch_size = temporal_patch_size

516

517

518 @dataclass(frozen=True, kw_only=True)

@@ -560,6 +597,7 @@ class RadioInternVisionModel(nn.Module):

597 max_img_size = int(

598 round(config.cpe_max_size / config.patch_size) * config.patch_size

599 )

600 + self.temporal_patch_size = config.video_temporal_patch_size

601 unique_teachers = set(t["name"] for t in config.teachers)

602 self.patch_generator = ViTPatchGenerator(

603 config.patch_size,

@@ -569,6 +607,8 @@ class RadioInternVisionModel(nn.Module):

607 cls_token=True,

608 num_cls_tokens=len(unique_teachers) if config.cls_token_per_teacher else 1,

609 register_multiple=config.register_multiple,

610 + temporal_patch_size=self.temporal_patch_size,

611 + separate_video_embedder=config.separate_video_embedder,

612 )

613

614 self.encoder = RadioVisionEncoder(

@@ -593,33 +633,68 @@ class RadioInternVisionModel(nn.Module):

633 def inter_image_mask_metadata(

634 self, imgs_sizes: list[tuple[int, int]], device: torch.device

635 ) -> MaskMetadata:

636 + """Build mask metadata from image pixel sizes. Adds num_skip to each

637 + sequence length (cls/register tokens) to match patch generator output."""

638 patch_size = self.patch_generator.patch_size

639 num_skip = self.patch_generator.num_skip

640

641 seq_lens = calc_seq_lens(imgs_sizes, patch_size)

642 adjusted = [s + num_skip for s in seq_lens]

643 + return self._inter_image_mask_metadata_from_seq_lens(adjusted, device=device)

644 +

645 + def _inter_image_mask_metadata_from_seq_lens(

646 + self, seq_lens: list[int], device: torch.device

647 + ) -> MaskMetadata:

648 + """Build mask metadata from actual sequence lengths (already including

649 + cls/register tokens, i.e. patch_count + num_skip per item).

650 + Use inter_image_mask_metadata() when you only have imgs_sizes."""

651 + assert len(seq_lens) > 0

652 cu_seqlens = torch.tensor(

602 - list(accumulate(adjusted, initial=0)), dtype=torch.int32, device=device

653 + list(accumulate(seq_lens, initial=0)), dtype=torch.int32, device=device

654 )

655 # Keep max_seqlen on CPU to avoid .item() sync

656 # See: https://github.com/vllm-project/vllm/blob/20b6b01/vllm/v1/attention/ops/vit_attn_wrappers.py#L48

606 - max_seqlen = torch.tensor(max(adjusted), dtype=torch.int32)

657 + max_seqlen = torch.tensor(max(seq_lens), dtype=torch.int32)

658 return MaskMetadata(cu_seqlens=cu_seqlens, max_seqlen=max_seqlen)

659

660 def forward(

661 self,

662 x: torch.Tensor,

663 imgs_sizes: list[tuple[int, int]] | None = None,

664 + num_frames: int | None = None,

665 ) -> torch.FloatTensor:

614 - hidden_states = self.patch_generator(x, imgs_sizes=imgs_sizes)

666 + T = self.temporal_patch_size

667 +

668 + # Build packed-sequence metadata for MMEncoderAttention when needed.

669 mask_meta = None

616 - if imgs_sizes is not None:

617 - assert len(imgs_sizes) > 0

618 - # Dynamic resolution: process each image as an independent sequence.

619 - mask_meta = self.inter_image_mask_metadata(

620 - imgs_sizes, device=hidden_states.device

670 + packed_batch_size = None # Original batch size before packing

671 +

672 + if num_frames is not None and T > 1:

673 + # Conv3d video: all tubelets have the same sequence length.

674 + # Pack [num_tubelets, seq_per_tubelet, hidden] → [1, total, hidden]

675 + hidden_states = self.patch_generator.forward_video(x)

676 + packed_batch_size, seq_per_tubelet, hidden_dim = hidden_states.shape

677 + hidden_states = hidden_states.reshape(1, -1, hidden_dim)

678 + mask_meta = self._inter_image_mask_metadata_from_seq_lens(

679 + [seq_per_tubelet] * packed_batch_size, device=hidden_states.device

680 )

681 + else:

682 + # Images for any model, or video for non-conv3d model

683 + hidden_states = self.patch_generator(x, imgs_sizes=imgs_sizes)

684 + if imgs_sizes is not None and len(imgs_sizes) > 1:

685 + # Dynamic resolution w/ > 1 image, create attn mask

686 + mask_meta = self.inter_image_mask_metadata(

687 + imgs_sizes, device=hidden_states.device

688 + )

689 +

690 encoder_outputs = self.encoder(inputs_embeds=hidden_states, mask_meta=mask_meta)

691 +

692 + # Unpack back to original batch shape if we packed for video

693 + if packed_batch_size is not None:

694 + encoder_outputs = encoder_outputs.reshape(

695 + packed_batch_size, seq_per_tubelet, -1

696 + )

697 +

698 return encoder_outputs

699

700

@@ -663,8 +738,13 @@ class RadioModel(nn.Module):

738 pixel_embeds: torch.Tensor | None = None,

739 *,

740 imgs_sizes: list[tuple[int, int]] | None = None,

741 + num_frames: int | None = None,

742 ) -> tuple[torch.FloatTensor, torch.FloatTensor]:

667 - y = self.model(pixel_values, imgs_sizes=imgs_sizes)

743 + y = self.model(

744 + pixel_values,

745 + imgs_sizes=imgs_sizes,

746 + num_frames=num_frames,

747 + )

748 return self._extract_final(y, imgs_sizes=imgs_sizes)

749

750 def load_weights(self, weights) -> set[str]:

@@ -714,6 +794,9 @@ class RadioModel(nn.Module):

794 weight_loader(param, weight)

795 loaded_params.add(vllm_key)

796

797 + if "model.patch_generator.video_embedder.weight" in loaded_params:

798 + self.model.patch_generator._video_embedder_loaded = True

799 +

800 return loaded_params

801

802 def _extract_final(

vllm/transformers_utils/configs/radio.py

+12

@@ -47,6 +47,14 @@ class RadioConfig(PretrainedConfig):

47 teachers: A list of teacher model configurations. Each teacher configuration is

48 a dict with keys like "name" and some may have "use_summary".

49 cls_token_per_teacher: Whether to use a separate CLS token for each teacher.

50 + video_temporal_patch_size: Number of consecutive video frames grouped into

51 + a single tubelet for temporal compression. Default 1 (no compression).

52 + When > 1, a dedicated video_embedder (3*T*P*P -> hidden) is created

53 + alongside the image embedder (3*P*P -> hidden).

54 + separate_video_embedder: When True and video_temporal_patch_size > 1, use a

55 + dedicated video patch embedder (3*T*P*P -> hidden) separate from the

56 + image embedder (3*P*P -> hidden). When False, a single embedder with

57 + input size 3*T*P*P is used for both (images are duplicated T times).

58 """

60 model_type = "radio"

@@ -68,6 +76,8 @@ class RadioConfig(PretrainedConfig):

76 register_multiple: int | None = None,

77 teachers: list[dict[str, Any]] | None = None,

78 cls_token_per_teacher: bool = False,

79 + video_temporal_patch_size: int = 1,

80 + separate_video_embedder: bool = True,

81 **kwargs,

82 ):

83 self.model_name = model_name

@@ -95,4 +105,6 @@ class RadioConfig(PretrainedConfig):

105 self.register_multiple = register_multiple

106 self.teachers = teachers if teachers is not None else []

107 self.cls_token_per_teacher = cls_token_per_teacher

108 + self.video_temporal_patch_size = video_temporal_patch_size

109 + self.separate_video_embedder = separate_video_embedder

110 super().__init__(**kwargs)

vllm/transformers_utils/processors/nano_nemotron_vl.py

+205 -27

@@ -11,6 +11,7 @@ import math

11 from abc import ABC, abstractmethod

12 from collections.abc import Sequence

13 from dataclasses import dataclass

14 +from functools import cached_property

15 from typing import Any, TypeVar

17 import einops

@@ -43,6 +44,12 @@ AUDIO_CONTEXT = "<so_embedding>"

44 # MAX_FRAMES = 16

45 DEFAULT_NUM_TILES = 12

47 +# Configure PIL to handle large images without warnings

48 +# This prevents DecompressionBombWarning for legitimate large images

49 +Image.MAX_IMAGE_PIXELS = None # Disable the limit entirely

50 +# Alternative: Set a specific higher limit

51 +# Image.MAX_IMAGE_PIXELS = 300000000 # ~300M pixels

52 +

54 def calculate_timestamps(

55 indices: list[int] | torch.Tensor,

@@ -138,19 +145,110 @@ def image_to_pixel_values(

145 return pixel_values

146

147

148 +def _compute_aspect_preserving_size(

149 + orig_w: int,

150 + orig_h: int,

151 + target_num_patches: int,

152 + patch_size: int,

153 + downsample_ratio: float,

154 +) -> tuple[int, int]:

155 + """Compute target pixel dimensions that preserve aspect ratio.

156 +

157 + Mirrors Megatron-LM image_processing.py video frame resizing:

158 + target area in patch-grid space is *target_num_patches*, distributed

159 + according to the source aspect ratio, then snapped to a multiple of

160 + the required divisor (2 for pixel-shuffle).

161 + """

162 + aspect_wh = orig_w / max(orig_h, 1)

163 + ph = round(math.sqrt(target_num_patches / aspect_wh))

164 + pw = round(math.sqrt(target_num_patches * aspect_wh))

165 + ph = max(ph, 1)

166 + pw = max(pw, 1)

167 +

168 + reduction_factor = int(round(1 / downsample_ratio))

169 + required_divisor = reduction_factor # 2 for pixel-shuffle

170 + if required_divisor > 1:

171 + rem_h = ph % required_divisor

172 + rem_w = pw % required_divisor

173 + ph_up = ph + (required_divisor - rem_h if rem_h else 0)

174 + ph_down = ph - rem_h

175 + pw_up = pw + (required_divisor - rem_w if rem_w else 0)

176 + pw_down = pw - rem_w

177 + if ph_up * pw_up <= target_num_patches:

178 + ph, pw = ph_up, pw_up

179 + else:

180 + ph = max(required_divisor, ph_down)

181 + pw = max(required_divisor, pw_down)

182 +

183 + return pw * patch_size, ph * patch_size # (width, height) in pixels

184 +

185 +

186 +def get_video_target_size_and_feature_size(

187 + orig_w: int,

188 + orig_h: int,

189 + target_patches: int,

190 + maintain_aspect_ratio: bool,

191 + patch_size: int,

192 + downsample_ratio: float,

193 +) -> tuple[int, int, int]:

194 + """Compute target (width, height) and feature_size for video resize and token count.

195 +

196 + Used by video_to_pixel_values (resize) and get_video_replacement_internvl

197 + (seq length calc) so both use the same dimensions.

198 + """

199 + if maintain_aspect_ratio:

200 + target_w, target_h = _compute_aspect_preserving_size(

201 + orig_w=orig_w,

202 + orig_h=orig_h,

203 + target_num_patches=target_patches,

204 + patch_size=patch_size,

205 + downsample_ratio=downsample_ratio,

206 + )

207 + else:

208 + reduction_factor = int(round(1 / downsample_ratio))

209 + side = int(math.sqrt(target_patches))

210 + side = max(reduction_factor, (side // reduction_factor) * reduction_factor)

211 + target_w = side * patch_size

212 + target_h = side * patch_size

213 +

214 + feature_size = int((target_h // patch_size) * downsample_ratio) * int(

215 + (target_w // patch_size) * downsample_ratio

216 + )

217 + return target_w, target_h, feature_size

218 +

219 +

220 def video_to_pixel_values(

221 video: npt.NDArray,

222 *,

223 input_size: int,

145 - max_num_tiles: int = 1,

146 - use_thumbnail: bool,

224 + video_target_num_patches: int | None = None,

225 + video_maintain_aspect_ratio: bool = False,

226 + patch_size: int = 16,

227 + downsample_ratio: float = 0.5,

228 ) -> torch.Tensor:

148 - assert max_num_tiles == 1, "Video modality always uses one tile"

149 -

229 # (num_frames, H, W, C) -> (num_frames, C, H, W)

230 video_tensor = torch.from_numpy(video).permute(0, 3, 1, 2)

231

153 - if video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size:

232 + if video_target_num_patches is not None:

233 + # Resize to target patch count (aspect-preserving or square).

234 + orig_h, orig_w = video_tensor.shape[2], video_tensor.shape[3]

235 + target_w, target_h, _ = get_video_target_size_and_feature_size(

236 + orig_w=orig_w,

237 + orig_h=orig_h,

238 + target_patches=video_target_num_patches,

239 + maintain_aspect_ratio=video_maintain_aspect_ratio,

240 + patch_size=patch_size,

241 + downsample_ratio=downsample_ratio,

242 + )

243 + if video_tensor.shape[2] != target_h or video_tensor.shape[3] != target_w:

244 + video_tensor = torch.nn.functional.interpolate(

245 + video_tensor,

246 + size=(target_h, target_w),

247 + mode="bicubic",

248 + align_corners=False,

249 + antialias=True,

250 + )

251 + elif video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size:

252 video_tensor = torch.nn.functional.interpolate(

253 video_tensor,

254 size=(input_size, input_size),

@@ -645,9 +743,9 @@ class BaseNanoNemotronVLProcessor(ABC):

743 "which should be a single string"

744 )

745 parts = [x for x in re.split(r"(<image>)", text[0]) if x]

648 - assert parts.count("<image>") == len(pixel_values_lst), (

649 - "the number of <image> tokens in the text should be the "

650 - "same as the number of images"

746 + assert parts.count("<image>") == len(num_tokens_per_image), (

747 + f"Expected {len(num_tokens_per_image)} <image> tokens in text "

748 + f"but found {parts.count('<image>')}"

749 )

750

751 for i, (feature_size, num_patches) in enumerate(

@@ -706,6 +804,33 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

804 self.video_token = video_token

805 self.video_pruning_rate = video_pruning_rate

806

807 + # Video params live exclusively in vision_config

808 + vision_config = getattr(config, "vision_config", config)

809 + self.video_temporal_patch_size: int = getattr(

810 + vision_config, "video_temporal_patch_size", 1

811 + )

812 + self.video_maintain_aspect_ratio: bool = getattr(

813 + vision_config, "video_maintain_aspect_ratio", False

814 + )

815 +

816 + # Resolve video frame target size: exactly one of video_target_num_patches

817 + # or video_target_img_size may be set (mirrors Megatron's

818 + # DynamicResolutionImageTilingStrategy validation).

819 + target_num_patches = getattr(vision_config, "video_target_num_patches", None)

820 + target_img_size = getattr(vision_config, "video_target_img_size", None)

821 + if target_num_patches is not None and target_img_size is not None:

822 + raise ValueError(

823 + "Exactly one of video_target_num_patches or "

824 + "video_target_img_size must be set, got both"

825 + )

826 + if target_num_patches is not None:

827 + self.video_target_num_patches: int | None = target_num_patches

828 + elif target_img_size is not None:

829 + base_patches = math.ceil(target_img_size / config.patch_size)

830 + self.video_target_num_patches = base_patches * base_patches

831 + else:

832 + self.video_target_num_patches = None

833 +

834 self.audio_extractor: ParakeetExtractor | None = None

835 raw_sound_config = getattr(config, "sound_config", None)

836 if raw_sound_config is not None:

@@ -721,6 +846,27 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

846 IMG_CONTEXT, add_special_tokens=False

847 )

848

849 + @cached_property

850 + def num_video_token(self) -> int:

851 + """Token count per video frame, accounting for video_target_num_patches.

852 +

853 + When video_target_num_patches is set the per-frame feature count

854 + differs from the image-based num_image_token. We use a square

855 + dummy (1:1) to compute the feature_size because the dummy video is

856 + square and the user confirmed that is acceptable.

857 + """

858 + if self.video_target_num_patches is not None:

859 + _, _, feature_size = get_video_target_size_and_feature_size(

860 + orig_w=self.image_size,

861 + orig_h=self.image_size,

862 + target_patches=self.video_target_num_patches,

863 + maintain_aspect_ratio=self.video_maintain_aspect_ratio,

864 + patch_size=self.config.patch_size,

865 + downsample_ratio=self.config.downsample_ratio,

866 + )

867 + return feature_size

868 + return self.num_image_token

869 +

870 @property

871 def supports_video(self) -> bool:

872 return self.video_token_id is not None

@@ -738,14 +884,15 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

884 def _videos_to_pixel_values_lst(

885 self,

886 videos: list[npt.NDArray],

741 - max_num_tiles: int,

887 ) -> list[torch.Tensor]:

888 return [

889 video_to_pixel_values(

890 video,

891 input_size=self.image_size,

747 - max_num_tiles=max_num_tiles,

748 - use_thumbnail=self.use_thumbnail,

892 + video_target_num_patches=self.video_target_num_patches,

893 + video_maintain_aspect_ratio=self.video_maintain_aspect_ratio,

894 + patch_size=self.config.patch_size,

895 + downsample_ratio=self.config.downsample_ratio,

896 )

897 for video in videos

898 ]

@@ -754,7 +901,6 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

901 self,

902 text: list[str],

903 videos: list[tuple[npt.NDArray, dict[str, Any]]],

757 - max_num_tiles: int,

904 ) -> tuple[list[str], dict[str, Any]]:

905 if len(videos) == 0 or not self.supports_video:

906 return text, {}

@@ -763,7 +909,6 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

909 video_metadata_lst = [v[1] for v in videos]

910 pixel_values_lst_video = self._videos_to_pixel_values_lst(

911 videos_lst,

766 - max_num_tiles=max_num_tiles,

912 )

913

914 # We use frame duration in milliseconds (as integer) to ensure

@@ -788,12 +933,10 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

933 "frame_duration_ms": torch.tensor(frame_duration_ms_lst),

934 }

935

791 - image_size: int = self.config.force_image_size

936 patch_size: int = self.config.patch_size

937 downsample_ratio = self.config.downsample_ratio

794 - tokens_in_single_frame = int(

795 - (image_size * image_size // patch_size**2) * (downsample_ratio**2)

796 - )

938 +

939 + T = self.video_temporal_patch_size

940

941 for pixel_values, video_metadata, frames_indices, frame_duration_ms in zip(

942 pixel_values_lst_video,

@@ -802,23 +945,28 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

945 frame_duration_ms_lst,

946 ):

947 num_frames = pixel_values.shape[0]

948 + frame_h, frame_w = pixel_values.shape[-2], pixel_values.shape[-1]

949 + tokens_in_single_frame = int(

950 + (frame_h * frame_w // patch_size**2) * (downsample_ratio**2)

951 + )

952 + num_tubelets = math.ceil(num_frames / T) if T > 1 else num_frames

953

954 if self.video_pruning_rate is not None and self.video_pruning_rate > 0.0:

955 # Start of EVS-specific code

956 num_tokens = compute_retained_tokens_count(

957 tokens_per_frame=tokens_in_single_frame,

810 - num_frames=num_frames,

958 + num_frames=num_tubelets,

959 q=self.video_pruning_rate,

960 )

961

962 # Here we just need placeholders that won't actually be replaced -

963 # we just need to make sure the total number of tokens is correct

964 # assign all tokens to the first frame

817 - tokens_per_frame = [num_tokens] + [0] * (num_frames - 1)

965 + tokens_per_frame = [num_tokens] + [0] * (num_tubelets - 1)

966

967 # End of EVS-specific code

968 else:

821 - tokens_per_frame = [tokens_in_single_frame] * num_frames

969 + tokens_per_frame = [tokens_in_single_frame] * num_tubelets

970

971 video_repl = self.get_video_repl(

972 tokens_per_frame=tokens_per_frame,

@@ -828,6 +976,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

976 img_start_token_ids=self._img_start_token_ids,

977 img_end_token_ids=self._img_end_token_ids,

978 img_context_token_ids=self._img_context_token_ids,

979 + video_temporal_patch_size=T,

980 )

981

982 # video_repl.full is a list of token IDs

@@ -908,7 +1057,6 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

1057 text, video_inputs = self._preprocess_video(

1058 text=text,

1059 videos=videos,

911 - max_num_tiles=1,

1060 )

1061

1062 text, audio_inputs = self._preprocess_audio(

@@ -962,6 +1110,7 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

1110 img_start_token_ids: list[int],

1111 img_end_token_ids: list[int],

1112 img_context_token_ids: list[int],

1113 + video_temporal_patch_size: int = 1,

1114 ) -> PromptUpdateDetails[list[int]]:

1115 """

1116 Build prompt replacement for a video.

@@ -981,31 +1130,60 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

1130 - EVS real (called from get_real_video_repl_for_evs) - different value per frame

1131 Args:

1132 tokens_per_frame (list[int]): number of tokens per frame

984 - frames_indices (list[int]): frame indices

1133 + (one per tubelet when T > 1)

1134 + frames_indices (list[int]): orig. frame indices

1135 + (one per frame, before tubelet subsampling)

1136 frame_duration_ms (int): duration of each frame in milliseconds

986 - tokenizer (HfTokenizer): tokenizer to use for tokenizing frame separators

1137 + tokenizer (TokenizerLike): tokenizer to use for tokenizing frame separators

1138 img_start_token_ids (list[int]): pre-tokenized IMG_START tokens

1139 img_end_token_ids (list[int]): pre-tokenized IMG_END tokens

1140 img_context_token_ids (list[int]): pre-tokenized IMG_CONTEXT tokens

1141 + video_temporal_patch_size (int): temporal patch size for videos

1142 """

1143 # TODO: Add support of frame_duration_ms to be None

1144 # At preprocessing step we should allow absent / metadata without

1145 # frames_indices field.

1146 timestamps_enabled = frame_duration_ms is not None

995 -

996 - if timestamps_enabled:

1147 + T = video_temporal_patch_size

1148 + num_frames = len(frames_indices)

1149 +

1150 + if T > 1 and timestamps_enabled:

1151 + all_timestamps = calculate_timestamps(frames_indices, frame_duration_ms)

1152 +

1153 + frame_separators = []

1154 + for group_idx, i in enumerate(range(0, num_frames, T)):

1155 + group_frames = []

1156 + for j in range(T): # Every frame in the group

1157 + frame_idx = i + j

1158 + if frame_idx < num_frames:

1159 + # Valid idx (haven't padded to mult. of T yet)

1160 + ts = all_timestamps[frame_idx]

1161 + frame_str = "Frame" if j == 0 else "frame"

1162 + group_frames.append(

1163 + f"{frame_str} {frame_idx + 1} sampled at {ts:.2f} seconds"

1164 + )

1165 + if group_frames:

1166 + # Join by `and` if there are >1 frame, otherwise no `and`

1167 + # Prepend \n to match training format (except first group)

1168 + sep = " and ".join(group_frames) + ": "

1169 + if group_idx > 0:

1170 + sep = "\n" + sep

1171 + frame_separators.append(sep)

1172 + elif timestamps_enabled:

1173 timestamps = calculate_timestamps(frames_indices, frame_duration_ms)

1174

1175 assert len(timestamps) == len(tokens_per_frame), (

1176 "timestamps and tokens_per_frame must have the same length"

1177 )

1178 frame_separators = [

1003 - f"Frame {i + 1} sampled at {timestamp:.2f} seconds: "

1179 + ("\n" if i > 0 else "")

1180 + + f"Frame {i + 1} sampled at {timestamp:.2f} seconds: "

1181 for i, timestamp in enumerate(timestamps)

1182 ]

1183 else:

1184 frame_separators = [

1008 - f"Frame {i + 1}: " for i, _ in enumerate(tokens_per_frame)

1185 + ("\n" if i > 0 else "") + f"Frame {i + 1}: "

1186 + for i, _ in enumerate(tokens_per_frame)

1187 ]

1188

1189 # Tokenize frame separator independently

vllm/v1/worker/gpu_model_runner.py

+13 -3

@@ -420,8 +420,9 @@ class GPUModelRunner(

420 self.is_multimodal_raw_input_only_model = (

421 model_config.is_multimodal_raw_input_only_model

422 )

423 - # This will be overridden in load_model()

423 + # These will be overridden in load_model()

424 self.is_multimodal_pruning_enabled = False

425 + self.requires_sequential_video_encoding = False

426 # Set to True after init_routed_experts_capturer() completes.

427 # Prevents routed experts code from running during profiling/dummy run.

428 self.routed_experts_initialized = False

@@ -2625,17 +2626,23 @@ class GPUModelRunner(

2626 ):

2627 batch_outputs: MultiModalEmbeddings

2628

2628 - # EVS-related change.

2629 + # EVS and dynamic res video related change.

2630 # (ekhvedchenia): Temporary hack to limit peak memory usage when

2631 # processing multimodal data. This solves the issue with scheduler

2632 # putting too many video samples into a single batch. Scheduler

2633 # uses pruned vision tokens count to compare it versus compute

2634 # budget which is incorrect (Either input media size or non-pruned

2635 # output vision tokens count should be considered)

2636 + # dynamic res video for nemotron temporarily uses this hack via

2637 + # requires_sequential_video_encoding

2638 + # because it doesn't yet support video batching.

2639 # TODO(ywang96): Fix memory profiling to take EVS into account and

2640 # remove this hack.

2641 if (

2638 - self.is_multimodal_pruning_enabled

2642 + (

2643 + self.is_multimodal_pruning_enabled

2644 + or self.requires_sequential_video_encoding

2645 + )

2646 and modality == "video"

2647 and num_items > 1

2648 ):

@@ -4609,6 +4616,9 @@ class GPUModelRunner(

4616 and mm_config is not None

4617 and mm_config.is_multimodal_pruning_enabled()

4618 )

4619 + self.requires_sequential_video_encoding = hasattr(

4620 + self.get_model(), "requires_sequential_video_encoding"

4621 + ) # Temporary hack for dynamic res video w/o support for bs>1 yet

4622

4623 if (

4624 is_mixture_of_experts(self.model)