053e7daa79 · tduncan/Nvidia-Optimization-Stack-VLLM

@@ -4,7 +4,7 @@

4 import itertools

5 import math

6 from collections.abc import Iterable, Mapping, Sequence

7 -from typing import Annotated, Literal

7 +from typing import Annotated, Any, Literal

8

9 import torch

10 import torch.nn as nn

@@ -49,6 +49,7 @@ from vllm.utils.tensor_schema import TensorSchema, TensorShape

49 from .interfaces import (

50 IsHybrid,

51 MultiModalEmbeddings,

52 + SupportsEncoderCudaGraph,

53 SupportsLoRA,

54 SupportsMultiModal,

55 SupportsPP,

@@ -63,6 +64,17 @@ from .utils import (

64 from .vision import is_vit_use_data_parallel

65

66

67 +def _pad_cumulative_seqlens_buffer(

68 + dst: torch.Tensor,

69 + src: torch.Tensor,

70 +) -> None:

71 + n = src.shape[0]

72 + dst.zero_()

73 + dst[:n].copy_(src)

74 + if n < dst.shape[0]:

75 + dst[n:] = src[-1]

76 +

77 +

78 class Lfm2VLImagePixelInputs(TensorSchema):

79 """

80 Dimensions:

@@ -558,13 +570,26 @@ class Lfm2VLMultiModalProjector(nn.Module):

570

571 if gather_idx_parts:

572 gather_idx = torch.cat(gather_idx_parts).to(device=device)

561 - gathered = vision_features_packed.index_select(0, gather_idx)

562 - unshuffled = gathered.reshape(-1, factor * factor * hidden_size)

573 + return self.forward_with_gather_idx(vision_features_packed, gather_idx)

574 else:

575 unshuffled = vision_features_packed.new_empty(

576 (0, factor * factor * hidden_size)

577 )

578

579 + return self.forward_from_unshuffled(unshuffled)

580 +

581 + def forward_with_gather_idx(

582 + self,

583 + vision_features_packed: torch.Tensor,

584 + gather_idx: torch.Tensor,

585 + ) -> torch.Tensor:

586 + hidden_size = vision_features_packed.shape[-1]

587 + factor = self.factor

588 + gathered = vision_features_packed.index_select(0, gather_idx)

589 + unshuffled = gathered.reshape(-1, factor * factor * hidden_size)

590 + return self.forward_from_unshuffled(unshuffled)

591 +

592 + def forward_from_unshuffled(self, unshuffled: torch.Tensor) -> torch.Tensor:

593 if self.projector_use_layernorm:

594 unshuffled = self.layer_norm(unshuffled)

595 hidden_states = self.linear_1(unshuffled)

@@ -579,7 +604,12 @@ class Lfm2VLMultiModalProjector(nn.Module):

604 dummy_inputs=Lfm2VLDummyInputsBuilder,

605 )

606 class Lfm2VLForConditionalGeneration(

582 - nn.Module, SupportsMultiModal, SupportsLoRA, SupportsPP, IsHybrid

607 + nn.Module,

608 + SupportsMultiModal,

609 + SupportsEncoderCudaGraph,

610 + SupportsLoRA,

611 + SupportsPP,

612 + IsHybrid,

613 ):

614 merge_by_field_config = True

615

@@ -645,6 +675,7 @@ class Lfm2VLForConditionalGeneration(

675

676 self.config = config

677 self.vllm_config = vllm_config

678 + self.model_config = vllm_config.model_config

679 self.multimodal_config = multimodal_config

680 self.use_data_parallel = multimodal_config.mm_encoder_tp_mode == "data"

681

@@ -697,7 +728,7 @@ class Lfm2VLForConditionalGeneration(

728 self,

729 pixel_values: torch.FloatTensor,

730 spatial_shapes: torch.Tensor,

700 - ) -> torch.Tensor:

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

732 assert spatial_shapes.device.type == "cpu", (

733 "Expected `spatial_shapes` on CPU to avoid device-to-host sync in "

734 "variable-length packing."

@@ -759,23 +790,13 @@ class Lfm2VLForConditionalGeneration(

790 )

791 vision_features_packed = image_outputs_packed[0]

792

762 - factor = self.multi_modal_projector.factor

763 - projected_lengths_list: list[int] = []

764 - for (height, width), length in zip(spatial_shapes_list, lengths_list):

765 - if length <= 0:

766 - projected_lengths_list.append(0)

767 - continue

768 - if height % factor != 0 or width % factor != 0:

769 - raise ValueError(

770 - "spatial_shapes must be divisible by downsample_factor: "

771 - f"got ({height}, {width}) with factor={factor}."

772 - )

773 - projected_lengths_list.append((height // factor) * (width // factor))

774 -

793 projected_packed = self.multi_modal_projector(

794 vision_features_packed=vision_features_packed,

795 spatial_shapes=spatial_shapes,

796 )

797 + projected_lengths_list = self._get_lfm2vl_tile_output_lengths(

798 + spatial_shapes_list

799 + )

800

801 image_features: list[torch.Tensor] = []

802 offset = 0

@@ -819,6 +840,383 @@ class Lfm2VLForConditionalGeneration(

840

841 return self._process_image_input(image_input)

842

843 + def get_encoder_cudagraph_config(self):

844 + from vllm.v1.worker.encoder_cudagraph_defs import (

845 + EncoderCudaGraphConfig,

846 + )

847 +

848 + return EncoderCudaGraphConfig(

849 + modalities=["image"],

850 + buffer_keys=[

851 + "pixel_values_packed",

852 + "pos_embeds",

853 + "cu_seqlens",

854 + "max_seqlen",

855 + "gather_idx",

856 + ],

857 + out_hidden_size=self.config.text_config.hidden_size,

858 + padding_logics={

859 + "cu_seqlens": _pad_cumulative_seqlens_buffer,

860 + },

861 + )

862 +

863 + def get_max_frames_per_video(self) -> int:

864 + return 0

865 +

866 + def get_encoder_cudagraph_budget_range(

867 + self,

868 + vllm_config: VllmConfig,

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

870 + min_budget = self._get_lfm2vl_min_image_tokens()

871 + max_budget = min(

872 + vllm_config.scheduler_config.max_num_batched_tokens,

873 + self.model_config.max_model_len,

874 + )

875 + return min_budget, max_budget

876 +

877 + def _get_spatial_shapes_list(

878 + self,

879 + spatial_shapes: torch.Tensor,

880 + ) -> list[list[int]]:

881 + assert spatial_shapes.device.type == "cpu", (

882 + "Expected `spatial_shapes` on CPU to avoid device-to-host sync in "

883 + "variable-length packing."

884 + )

885 + return spatial_shapes.tolist()

886 +

887 + @staticmethod

888 + def _get_lfm2vl_tile_input_lengths(

889 + spatial_shapes_list: list[list[int]],

890 + ) -> list[int]:

891 + return [height * width for height, width in spatial_shapes_list]

892 +

893 + def _get_lfm2vl_tile_output_lengths(

894 + self,

895 + spatial_shapes_list: list[list[int]],

896 + ) -> list[int]:

897 + factor = self.multi_modal_projector.factor

898 + output_lengths: list[int] = []

899 + for height, width in spatial_shapes_list:

900 + if height % factor != 0 or width % factor != 0:

901 + raise ValueError(

902 + "spatial_shapes must be divisible by downsample_factor: "

903 + f"got ({height}, {width}) with factor={factor}."

904 + )

905 + output_lengths.append((height // factor) * (width // factor))

906 + return output_lengths

907 +

908 + def _get_lfm2vl_mm_processor_kwargs(self) -> Mapping[str, object]:

909 + return self.multimodal_config.mm_processor_kwargs or {}

910 +

911 + def _get_lfm2vl_min_image_tokens(self) -> int:

912 + value = self._get_lfm2vl_mm_processor_kwargs().get(

913 + "min_image_tokens",

914 + getattr(self.config, "min_image_tokens", None) or 64,

915 + )

916 + return max(1, int(value))

917 +

918 + def _get_lfm2vl_item_tile_slices(

919 + self,

920 + num_patches: torch.Tensor,

921 + ) -> list[tuple[int, int]]:

922 + num_patches_list = [int(x) for x in num_patches.tolist()]

923 + starts = [0]

924 + for count in num_patches_list:

925 + starts.append(starts[-1] + count)

926 + return list(zip(starts[:-1], starts[1:]))

927 +

928 + def get_encoder_cudagraph_item_specs(

929 + self,

930 + mm_kwargs: dict[str, Any],

931 + ):

932 + from vllm.v1.worker.encoder_cudagraph_defs import EncoderItemSpec

933 +

934 + spatial_shapes = mm_kwargs["spatial_shapes"]

935 + num_patches = mm_kwargs["num_patches"]

936 + spatial_shapes_list = self._get_spatial_shapes_list(spatial_shapes)

937 + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list)

938 + output_lengths = self._get_lfm2vl_tile_output_lengths(spatial_shapes_list)

939 +

940 + return [

941 + EncoderItemSpec(

942 + input_size=sum(input_lengths[start:end]),

943 + output_tokens=sum(output_lengths[start:end]),

944 + )

945 + for start, end in self._get_lfm2vl_item_tile_slices(num_patches)

946 + ]

947 +

948 + def select_encoder_cudagraph_items(

949 + self,

950 + mm_kwargs: dict[str, Any],

951 + indices: list[int],

952 + ) -> dict[str, Any]:

953 + pixel_values = mm_kwargs["pixel_values"]

954 + spatial_shapes = mm_kwargs["spatial_shapes"]

955 + num_patches = mm_kwargs["num_patches"]

956 +

957 + tile_slices = self._get_lfm2vl_item_tile_slices(num_patches)

958 +

959 + if len(indices) == 0:

960 + return {

961 + "pixel_values": pixel_values[:0],

962 + "spatial_shapes": spatial_shapes[:0],

963 + "num_patches": num_patches[:0],

964 + }

965 +

966 + tile_indices: list[int] = []

967 + for image_idx in indices:

968 + start, end = tile_slices[image_idx]

969 + tile_indices.extend(range(start, end))

970 +

971 + return {

972 + "pixel_values": pixel_values[tile_indices],

973 + "spatial_shapes": spatial_shapes[tile_indices],

974 + "num_patches": num_patches[indices],

975 + }

976 +

977 + def _pack_lfm2vl_pixel_values(

978 + self,

979 + pixel_values: torch.Tensor,

980 + spatial_shapes_list: list[list[int]],

981 + ) -> torch.Tensor:

982 + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list)

983 + total_tokens = sum(input_lengths)

984 + packed = pixel_values.new_empty((total_tokens, pixel_values.shape[-1]))

985 +

986 + offset = 0

987 + for i, length in enumerate(input_lengths):

988 + if length <= 0:

989 + continue

990 + packed[offset : offset + length].copy_(pixel_values[i, :length])

991 + offset += length

992 + return packed

993 +

994 + def _get_lfm2vl_pos_embeds(

995 + self,

996 + spatial_shapes: torch.Tensor,

997 + spatial_shapes_list: list[list[int]],

998 + ) -> torch.Tensor:

999 + embeddings = self.vision_tower.vision_model.embeddings

1000 + positional_embeddings = embeddings.position_embedding.weight.reshape(

1001 + embeddings.position_embedding_size,

1002 + embeddings.position_embedding_size,

1003 + -1,

1004 + )

1005 + lengths_list = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list)

1006 + return embeddings.resize_positional_embeddings_packed(

1007 + positional_embeddings,

1008 + spatial_shapes,

1009 + lengths_list=lengths_list,

1010 + )

1011 +

1012 + def _get_lfm2vl_cu_seqlens(

1013 + self,

1014 + spatial_shapes_list: list[list[int]],

1015 + device: torch.device,

1016 + ) -> torch.Tensor:

1017 + lengths = torch.tensor(

1018 + self._get_lfm2vl_tile_input_lengths(spatial_shapes_list),

1019 + dtype=torch.int32,

1020 + device=device,

1021 + )

1022 + cu_seqlens = torch.zeros(

1023 + lengths.shape[0] + 1,

1024 + dtype=torch.int32,

1025 + device=device,

1026 + )

1027 + if lengths.numel() > 0:

1028 + cu_seqlens[1:] = torch.cumsum(lengths, dim=0)

1029 + return cu_seqlens

1030 +

1031 + def _get_lfm2vl_max_seqlen(

1032 + self,

1033 + spatial_shapes_list: list[list[int]],

1034 + ) -> torch.Tensor:

1035 + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list)

1036 + max_seqlen = max(input_lengths) if input_lengths else 0

1037 + return torch.tensor(max_seqlen, dtype=torch.int32)

1038 +

1039 + def _get_lfm2vl_projector_gather_idx(

1040 + self,

1041 + spatial_shapes_list: list[list[int]],

1042 + device: torch.device,

1043 + ) -> torch.Tensor:

1044 + factor = self.multi_modal_projector.factor

1045 + dh = torch.arange(factor, dtype=torch.int64)

1046 + dw = torch.arange(factor, dtype=torch.int64)

1047 + dh_grid, dw_grid = torch.meshgrid(dh, dw, indexing="ij")

1048 + dh_flat = dh_grid.reshape(-1)

1049 + dw_flat = dw_grid.reshape(-1)

1050 +

1051 + gather_idx_parts: list[torch.Tensor] = []

1052 + offset = 0

1053 + for height, width in spatial_shapes_list:

1054 + length = height * width

1055 + if length <= 0:

1056 + continue

1057 + if height % factor != 0 or width % factor != 0:

1058 + raise ValueError(

1059 + "spatial_shapes must be divisible by downsample_factor: "

1060 + f"got ({height}, {width}) with factor={factor}."

1061 + )

1062 +

1063 + rows_out = torch.arange(height // factor, dtype=torch.int64)

1064 + cols_out = torch.arange(width // factor, dtype=torch.int64)

1065 + rr, cc = torch.meshgrid(rows_out, cols_out, indexing="ij")

1066 + rr = rr.reshape(-1)

1067 + cc = cc.reshape(-1)

1068 + token_idx = (rr[:, None] * factor + dh_flat[None, :]) * width + (

1069 + cc[:, None] * factor + dw_flat[None, :]

1070 + )

1071 + gather_idx_parts.append(token_idx.reshape(-1) + offset)

1072 + offset += length

1073 +

1074 + if not gather_idx_parts:

1075 + return torch.empty(0, dtype=torch.int64, device=device)

1076 + return torch.cat(gather_idx_parts).to(device=device)

1077 +

1078 + def _prepare_lfm2vl_cudagraph_values(

1079 + self,

1080 + pixel_values: torch.Tensor,

1081 + spatial_shapes: torch.Tensor,

1082 + ) -> dict[str, torch.Tensor]:

1083 + spatial_shapes_list = self._get_spatial_shapes_list(spatial_shapes)

1084 + pixel_values_packed = self._pack_lfm2vl_pixel_values(

1085 + pixel_values,

1086 + spatial_shapes_list,

1087 + )

1088 + pos_embeds = self._get_lfm2vl_pos_embeds(spatial_shapes, spatial_shapes_list)

1089 + device = pixel_values.device

1090 +

1091 + return {

1092 + "pixel_values_packed": pixel_values_packed,

1093 + "pos_embeds": pos_embeds,

1094 + "cu_seqlens": self._get_lfm2vl_cu_seqlens(spatial_shapes_list, device),

1095 + "max_seqlen": self._get_lfm2vl_max_seqlen(spatial_shapes_list),

1096 + "gather_idx": self._get_lfm2vl_projector_gather_idx(

1097 + spatial_shapes_list,

1098 + device,

1099 + ),

1100 + }

1101 +

1102 + def _get_lfm2vl_capture_spatial_shapes(

1103 + self,

1104 + token_budget: int,

1105 + ) -> torch.Tensor:

1106 + factor = self.multi_modal_projector.factor

1107 + min_image_tokens = self._get_lfm2vl_min_image_tokens()

1108 + remaining = token_budget

1109 + shapes: list[list[int]] = []

1110 +

1111 + while remaining > 0:

1112 + out_tokens = min(remaining, min_image_tokens)

1113 + shapes.append([factor, out_tokens * factor])

1114 + remaining -= out_tokens

1115 +

1116 + return torch.tensor(shapes, dtype=torch.int64)

1117 +

1118 + def prepare_encoder_cudagraph_capture_inputs(

1119 + self,

1120 + token_budget: int,

1121 + max_batch_size: int,

1122 + max_frames_per_batch: int,

1123 + device: torch.device,

1124 + dtype: torch.dtype,

1125 + ):

1126 + from vllm.v1.worker.encoder_cudagraph_defs import (

1127 + EncoderCudaGraphCaptureInputs,

1128 + )

1129 +

1130 + spatial_shapes = self._get_lfm2vl_capture_spatial_shapes(token_budget)

1131 + spatial_shapes_list = self._get_spatial_shapes_list(spatial_shapes)

1132 + input_lengths = self._get_lfm2vl_tile_input_lengths(spatial_shapes_list)

1133 + total_input_tokens = sum(input_lengths)

1134 +

1135 + patch_dim = (

1136 + self.vision_tower.vision_model.embeddings.patch_embedding.weight.shape[1]

1137 + )

1138 + dummy_pixel_values = torch.randn(

1139 + total_input_tokens,

1140 + patch_dim,

1141 + device=device,

1142 + dtype=dtype,

1143 + )

1144 + pos_embeds = self._get_lfm2vl_pos_embeds(

1145 + spatial_shapes,

1146 + spatial_shapes_list,

1147 + ).to(device=device, dtype=dtype)

1148 +

1149 + # max_seqlen.item() is baked into the captured ViT attention graph, so

1150 + # capture with a budget-level upper bound that covers any replay item.

1151 + max_tile_input_tokens = token_budget * self.multi_modal_projector.factor**2

1152 + values = {

1153 + "pixel_values_packed": dummy_pixel_values,

1154 + "pos_embeds": pos_embeds,

1155 + "cu_seqlens": self._get_lfm2vl_cu_seqlens(spatial_shapes_list, device),

1156 + "max_seqlen": torch.tensor(max_tile_input_tokens, dtype=torch.int32),

1157 + "gather_idx": self._get_lfm2vl_projector_gather_idx(

1158 + spatial_shapes_list,

1159 + device,

1160 + ),

1161 + }

1162 +

1163 + return EncoderCudaGraphCaptureInputs(values=values)

1164 +

1165 + def prepare_encoder_cudagraph_replay_buffers(

1166 + self,

1167 + mm_kwargs: dict[str, Any],

1168 + max_batch_size: int,

1169 + max_frames_per_batch: int,

1170 + ):

1171 + from vllm.v1.worker.encoder_cudagraph_defs import (

1172 + EncoderCudaGraphReplayBuffers,

1173 + )

1174 +

1175 + values = self._prepare_lfm2vl_cudagraph_values(

1176 + mm_kwargs["pixel_values"],

1177 + mm_kwargs["spatial_shapes"],

1178 + )

1179 + return EncoderCudaGraphReplayBuffers(values=values)

1180 +

1181 + def encoder_cudagraph_forward(

1182 + self,

1183 + values: dict[str, torch.Tensor],

1184 + ) -> torch.Tensor:

1185 + embeddings = self.vision_tower.vision_model.embeddings

1186 + pixel_values = values["pixel_values_packed"].to(

1187 + dtype=embeddings.patch_embedding.weight.dtype

1188 + )

1189 + patch_embeds = embeddings.patch_embedding(pixel_values)

1190 + hidden_states = (patch_embeds + values["pos_embeds"]).unsqueeze(0)

1191 +

1192 + with set_forward_context(None, self.vllm_config):

1193 + encoder_outputs = self.vision_tower.vision_model.encoder(

1194 + inputs_embeds=hidden_states,

1195 + cu_seqlens=values["cu_seqlens"],

1196 + max_seqlen=values["max_seqlen"],

1197 + )

1198 +

1199 + post_layernorm = self.vision_tower.vision_model.post_layernorm

1200 + if post_layernorm is not None:

1201 + encoder_outputs = post_layernorm(encoder_outputs)

1202 +

1203 + return self.multi_modal_projector.forward_with_gather_idx(

1204 + vision_features_packed=encoder_outputs[0],

1205 + gather_idx=values["gather_idx"],

1206 + )

1207 +

1208 + def encoder_eager_forward(

1209 + self,

1210 + mm_kwargs: dict[str, Any],

1211 + ) -> torch.Tensor:

1212 + image_input = LFM2VLImageInputs(

1213 + type="pixel_values",

1214 + pixel_values=mm_kwargs["pixel_values"],

1215 + spatial_shapes=mm_kwargs["spatial_shapes"],

1216 + num_patches=mm_kwargs["num_patches"],

1217 + )

1218 + return torch.cat(self._process_image_input(image_input), dim=0)

1219 +

1220 def forward(

1221 self,

1222 input_ids: torch.Tensor | None,

[Model] Add encoder CUDA graph support to Lfm2VL (#44930)