982beae809 · tduncan/Nvidia-Optimization-Stack-VLLM

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

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

9

10 import math

11 -import warnings

11 from abc import ABC, abstractmethod

12 from collections.abc import Sequence

13 from dataclasses import dataclass

@@ -26,7 +25,7 @@ from transformers import BatchFeature, PretrainedConfig, TensorType

25 from vllm.model_executor.models.parakeet import ParakeetExtractor

26 from vllm.multimodal.evs import compute_retained_tokens_count

27 from vllm.multimodal.inputs import AudioItem

29 -from vllm.multimodal.processing.processor import PromptUpdateDetails, _seq2tokens

28 +from vllm.multimodal.processing.processor import PromptUpdateDetails

29 from vllm.tokenizers.hf import HfTokenizer

30

31 from .internvl import calculate_internvl_targets, get_internvl_target_ratios

@@ -63,42 +62,50 @@ def calculate_timestamps(

62 return timestamps

63

64

66 -def input_conditioner(x: torch.Tensor, norm_mean: torch.Tensor, norm_std: torch.Tensor):

67 - return (x - norm_mean) / norm_std

65 +@torch.compile(dynamic=True)

66 +def _bicubic_resize_and_normalize(

67 + tensor: torch.Tensor,

68 + size: tuple[int, int] | None = None,

69 + norm_mean: torch.Tensor | None = None,

70 + norm_std: torch.Tensor | None = None,

71 + dtype: torch.dtype = torch.float32,

72 +) -> torch.Tensor:

73 + """Permute NHWC→NCHW, optional bicubic resize, rescale + normalize.

74

75 + Input must be a raw 4-D **NHWC** tensor.

76

70 -def _bicubic_from_ndarray(

71 - array: npt.NDArray[Any], *, size: tuple[int, int]

72 -) -> torch.Tensor:

73 - """

74 - Convert a 4D NHWC ndarray to NCHW and interpolate with bicubic.

75 - Suppresses PyTorch's non-writable NumPy warning because interpolate copies,

76 - and torch.from_numpy(array) is discarded at the end of function scope.

77 + *size*: target ``(H, W)``; skips interpolation when ``None``.

78 + *norm_mean* / *norm_std*: when both provided, fused

79 + ``(x/255 - mean) / std`` + dtype cast; otherwise ``x/255`` + cast.

80 """

78 -

79 - with warnings.catch_warnings():

80 - msg = "The given NumPy array is not writ.*"

81 - # Apparently, different versions of PyTorch use writable or writeable.

82 - warnings.filterwarnings("ignore", message=msg, category=UserWarning)

83 - tensor = torch.from_numpy(array)

84 - assert tensor.ndim == 4, f"{tensor.ndim=}"

85 - tensor = tensor.permute(0, 3, 1, 2)

86 - return (

87 - torch.nn.functional.interpolate(

81 + tensor = tensor.permute(0, 3, 1, 2).to(dtype=torch.float32)

82 + if size is not None:

83 + tensor = torch.nn.functional.interpolate(

84 tensor, size=size, mode="bicubic", align_corners=False, antialias=True

85 )

90 - / 255.0

86 + if norm_mean is not None and norm_std is not None:

87 + return ((tensor / 255.0 - norm_mean) / norm_std).to(dtype=dtype).contiguous()

88 + return (tensor / 255.0).to(dtype=dtype).contiguous()

89 +

90 +

91 +def _pil_to_nhwc_tensor(image: Image.Image) -> torch.Tensor:

92 + """Convert a PIL image to a 4-D NHWC tensor suitable for compiled ops."""

93 + array = np.asarray(

94 + image.convert("RGB") if image.mode != "RGB" else image, dtype=np.uint8

95 )

96 + return torch.from_numpy(np.expand_dims(array, axis=0))

97

98

99 def dynamic_preprocess(

95 - image,

100 + image: Image.Image,

101 *,

97 - image_size=512,

98 - max_num_tiles=12,

99 - use_thumbnail=True,

100 - idx=0,

101 -):

102 + image_size: int = 512,

103 + max_num_tiles: int = 12,

104 + use_thumbnail: bool = True,

105 + norm_mean: torch.Tensor | None = None,

106 + norm_std: torch.Tensor | None = None,

107 + dtype: torch.dtype = torch.float32,

108 +) -> torch.Tensor:

109 orig_width, orig_height = image.size

110

111 target_ratios = get_internvl_target_ratios(1, max_num_tiles)

@@ -111,13 +118,15 @@ def dynamic_preprocess(

118 use_thumbnail=False,

119 )

120

114 - image = np.asarray(

115 - image.convert("RGB") if image.mode != "RGB" else image, dtype=np.uint8

116 - )

121 + tensor = _pil_to_nhwc_tensor(image)

122

118 - image = np.expand_dims(image, axis=0)

119 -

120 - resized_img = _bicubic_from_ndarray(image, size=(target_height, target_width))

123 + resized_img = _bicubic_resize_and_normalize(

124 + tensor,

125 + size=(target_height, target_width),

126 + norm_mean=norm_mean,

127 + norm_std=norm_std,

128 + dtype=dtype,

129 + )

130 B, C, H, W = resized_img.shape

131 hp, wp = H // image_size, W // image_size

132 patches = (

@@ -127,30 +136,16 @@ def dynamic_preprocess(

136 )

137

138 if use_thumbnail and patches.shape[0] > 1:

130 - thumb = _bicubic_from_ndarray(image, size=(image_size, image_size))

139 + thumb = _bicubic_resize_and_normalize(

140 + tensor,

141 + size=(image_size, image_size),

142 + norm_mean=norm_mean,

143 + norm_std=norm_std,

144 + dtype=dtype,

145 + )

146 patches = torch.cat([patches, thumb], dim=0)

147

133 - return list(patches)

134 -

135 -

136 -def image_to_pixel_values(

137 - image: Image.Image,

138 - *,

139 - input_size: int,

140 - max_num: int,

141 - use_thumbnail: bool,

142 - idx: int,

143 -) -> torch.Tensor:

144 - images = dynamic_preprocess(

145 - image,

146 - image_size=input_size,

147 - max_num_tiles=max_num,

148 - use_thumbnail=use_thumbnail,

149 - idx=idx,

150 - )

151 -

152 - pixel_values = torch.stack(images)

153 - return pixel_values

148 + return patches

149

150

151 def _compute_aspect_preserving_size(

@@ -233,14 +228,16 @@ def video_to_pixel_values(

228 video_maintain_aspect_ratio: bool = False,

229 patch_size: int = 16,

230 downsample_ratio: float = 0.5,

231 + norm_mean: torch.Tensor | None = None,

232 + norm_std: torch.Tensor | None = None,

233 + dtype: torch.dtype = torch.float32,

234 ) -> torch.Tensor:

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

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

235 + """Convert video ndarray (T, H, W, C) to normalized pixel tensor (T, C, H, W)."""

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

237 + size: tuple[int, int] | None = None

238

239 if video_target_num_patches is not None:

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

242 - orig_h, orig_w = video_tensor.shape[2], video_tensor.shape[3]

243 - target_w, target_h, _ = get_video_target_size_and_feature_size(

240 + tw, th, _ = get_video_target_size_and_feature_size(

241 orig_w=orig_w,

242 orig_h=orig_h,

243 target_patches=video_target_num_patches,

@@ -248,14 +245,13 @@ def video_to_pixel_values(

245 patch_size=patch_size,

246 downsample_ratio=downsample_ratio,

247 )

251 - if video_tensor.shape[2] != target_h or video_tensor.shape[3] != target_w:

252 - return _bicubic_from_ndarray(video, size=(target_h, target_w))

253 - elif video_tensor.shape[2] != input_size or video_tensor.shape[3] != input_size:

254 - return _bicubic_from_ndarray(video, size=(input_size, input_size))

248 + if orig_h != th or orig_w != tw:

249 + size = (th, tw)

250 + elif orig_h != input_size or orig_w != input_size:

251 + size = (input_size, input_size)

252

256 - video_tensor = video_tensor / 255.0

257 -

258 - return video_tensor

253 + tensor = torch.from_numpy(video)

254 + return _bicubic_resize_and_normalize(tensor, size, norm_mean, norm_std, dtype)

255

256

257 class DynamicResolutionImageTiler:

@@ -343,6 +339,7 @@ class DynamicResolutionImageTiler:

339 self,

340 text_prompt_length: int,

341 images: list[Image.Image],

342 + dtype: torch.dtype = torch.float32,

343 ) -> tuple[list[torch.Tensor], list[int]]:

344 num_tokens_available = self.max_num_tokens_available(text_prompt_length)

345 params_per_image = self.compute_params(images, num_tokens_available)

@@ -350,7 +347,7 @@ class DynamicResolutionImageTiler:

347 feature_sizes = []

348 images = []

349 for param in params_per_image:

353 - for t in self.apply_params(param):

350 + for t in self.apply_params(param, dtype=dtype):

351 assert t.ndim == 3, f"{t.ndim=}: expected 3 dim tensor"

352 images.append(t)

353 feature_sizes.append(param.num_embeddings)

@@ -363,17 +360,23 @@ class DynamicResolutionImageTiler:

360 num_embeddings: int

361 patch_size: tuple[int, int]

362

366 - def apply_params(self, params: DynamicResolutionParams) -> list[torch.Tensor]:

363 + def apply_params(

364 + self,

365 + params: DynamicResolutionParams,

366 + dtype: torch.dtype = torch.float32,

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

368 target_size = (

369 params.patch_size[1] * self._patch_size,

370 params.patch_size[0] * self._patch_size,

371 )

371 - image = np.asarray(

372 - params.media.convert("RGB") if params.media.mode != "RGB" else params.media,

373 - dtype=np.uint8,

372 + tensor = _pil_to_nhwc_tensor(params.media)

373 + resized_img = _bicubic_resize_and_normalize(

374 + tensor,

375 + size=target_size,

376 + norm_mean=self.norm_mean,

377 + norm_std=self.norm_std,

378 + dtype=dtype,

379 )

375 - image = np.expand_dims(image, axis=0)

376 - resized_img = _bicubic_from_ndarray(image, size=target_size)

380 return list(resized_img)

381

382 def process_media(

@@ -619,6 +622,7 @@ class BaseNanoNemotronVLProcessor(ABC):

622 norm_mean=config.norm_mean,

623 norm_std=config.norm_std,

624 )

625 + self.dtype: torch.dtype = getattr(config, "dtype", torch.float32)

626

627 @staticmethod

628 def use_dynamic_resolution(config: PretrainedConfig) -> bool:

@@ -662,14 +666,16 @@ class BaseNanoNemotronVLProcessor(ABC):

666 max_num_tiles: int,

667 ) -> list[torch.Tensor]:

668 return [

665 - image_to_pixel_values(

669 + dynamic_preprocess(

670 image,

667 - input_size=self.image_size,

668 - max_num=max_num_tiles,

671 + image_size=self.image_size,

672 + max_num_tiles=max_num_tiles,

673 use_thumbnail=self.use_thumbnail,

670 - idx=idx,

674 + norm_mean=self.norm_mean,

675 + norm_std=self.norm_std,

676 + dtype=self.dtype,

677 )

672 - for idx, image in enumerate(images)

678 + for image in images

679 ]

680

681 def _preprocess_image(

@@ -690,23 +696,22 @@ class BaseNanoNemotronVLProcessor(ABC):

696 pixel_values_lst, num_tokens_per_image = tiler._images_to_pixel_values_lst(

697 text_prompt_length=text_prompt_length,

698 images=images,

699 + dtype=self.dtype,

700 )

701 imgs_sizes = [(pv.shape[-2], pv.shape[-1]) for pv in pixel_values_lst]

695 - normalized = [

696 - input_conditioner(img, tiler.norm_mean, tiler.norm_std)

697 - for img in pixel_values_lst

698 - ]

702 image_num_patches = torch.tensor([1] * len(num_tokens_per_image))

703 image_inputs = {

701 - "pixel_values_flat": normalized,

704 + "pixel_values_flat": pixel_values_lst,

705 "imgs_sizes": imgs_sizes,

706 "num_tokens_per_image": num_tokens_per_image,

707 }

708 else:

709 pixel_values_lst = self._images_to_pixel_values_lst(images, max_num_tiles)

710 image_num_patches = torch.tensor([len(item) for item in pixel_values_lst])

708 - pixel_values_flat = input_conditioner(

709 - torch.cat(pixel_values_lst), self.norm_mean, self.norm_std

711 + pixel_values_flat = (

712 + torch.cat(pixel_values_lst)

713 + if len(pixel_values_lst) > 1

714 + else pixel_values_lst[0]

715 )

716 image_inputs = {

717 "pixel_values_flat": pixel_values_flat,

@@ -863,6 +868,8 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

868 def _videos_to_pixel_values_lst(

869 self,

870 videos: list[npt.NDArray],

871 + *,

872 + dtype: torch.dtype = torch.float32,

873 ) -> list[torch.Tensor]:

874 return [

875 video_to_pixel_values(

@@ -872,6 +879,9 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

879 video_maintain_aspect_ratio=self.video_maintain_aspect_ratio,

880 patch_size=self.config.patch_size,

881 downsample_ratio=self.config.downsample_ratio,

882 + norm_mean=self.norm_mean,

883 + norm_std=self.norm_std,

884 + dtype=dtype,

885 )

886 for video in videos

887 ]

@@ -886,8 +896,10 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

896

897 videos_lst = [v[0] for v in videos]

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

899 +

900 pixel_values_lst_video = self._videos_to_pixel_values_lst(

901 videos_lst,

902 + dtype=self.dtype,

903 )

904

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

@@ -903,10 +915,15 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

915 metadata["frames_indices"] for metadata in video_metadata_lst

916 ]

917 video_num_patches = torch.tensor([len(item) for item in pixel_values_lst_video])

918 +

919 + # Normalization already fused into resize above.

920 + # Skip the torch.cat copy when there is exactly one video

921 + if len(pixel_values_lst_video) == 1:

922 + pixel_values_flat = pixel_values_lst_video[0]

923 + else:

924 + pixel_values_flat = torch.cat(pixel_values_lst_video)

925 video_inputs = {

907 - "pixel_values_flat_video": input_conditioner(

908 - torch.cat(pixel_values_lst_video), self.norm_mean, self.norm_std

909 - ),

926 + "pixel_values_flat_video": pixel_values_flat,

927 "video_num_patches": video_num_patches,

928 "frames_indices": frames_indices_lst,

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

@@ -1168,20 +1185,21 @@ class NanoNemotronVLProcessor(BaseNanoNemotronVLProcessor):

1185 for i, _ in enumerate(tokens_per_frame)

1186 ]

1187

1171 - # Tokenize frame separator independently

1172 - frame_separators_tokenized = [

1173 - _seq2tokens(tokenizer, sep) for sep in frame_separators

1174 - ]

1188 + # Batch-tokenize all frame separators at once — the HuggingFace

1189 + # tokenizers Rust backend parallelizes batch encoding across threads.

1190 + batch_encoded = tokenizer(

1191 + frame_separators,

1192 + add_special_tokens=False,

1193 + return_attention_mask=False,

1194 + )

1195 + frame_separators_tokenized: list[list[int]] = batch_encoded["input_ids"]

1196

1197 # Tokenize each component independently to avoid tokenizer merging tokens

1198 # across boundaries. This ensures consistent tokenization regardless of

1199 # num_tokens_per_frame values.

1200 all_token_ids = []

1201 for i, num_tokens in enumerate(tokens_per_frame):

1181 - frame_sep_token_ids = frame_separators_tokenized[i]

1182 - all_token_ids.extend(frame_sep_token_ids)

1183 -

1184 - # Add pre-tokenized special tokens

1202 + all_token_ids.extend(frame_separators_tokenized[i])

1203 all_token_ids.extend(img_start_token_ids)

1204 all_token_ids.extend(img_context_token_ids * num_tokens)

1205 all_token_ids.extend(img_end_token_ids)

Optimize nemotron VL image/video preprocessing (#40283)