@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / ddaff2938e
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 14, 05:30 UTC

[MM] Move Qwen3Omni MRoPE impl to model file (#26608)

Signed-off-by: Roger Wang <hey@rogerw.io> Co-authored-by: gemini-code-assist[bot] <176961590+gemini-code-assist[bot]@users.noreply.github.com>

Roger Wang committed Oct 10, 2025 at 22:17 UTC ddaff2938e0b78b2d6237f6f7975ac19167cb04e
4 files changed +368 -387
vllm/model_executor/layers/rotary_embedding/mrope.py
+1 -360
@@ -426,7 +426,7 @@ class MRotaryEmbedding(RotaryEmbedding):
426 ) -> tuple[torch.Tensor, int]:
427 from vllm.transformers_utils.config import thinker_uses_mrope
428
429 - if thinker_uses_mrope(hf_config):
429 + if thinker_uses_mrope(hf_config) and hf_config.model_type == "qwen2_5_omni":
430 return cls._omni_get_input_positions_tensor(
431 input_tokens=input_tokens,
432 hf_config=hf_config,
@@ -1119,339 +1119,6 @@ class MRotaryEmbedding(RotaryEmbedding):
1119
1120 return llm_positions, mrope_position_delta
1121
1122 - @classmethod
1123 - def _omni3_get_input_positions_tensor(
1124 - cls,
1125 - config,
1126 - input_ids: torch.Tensor,
1127 - image_grid_thw: torch.Tensor,
1128 - video_grid_thw: torch.Tensor,
1129 - use_audio_in_video: bool = False,
1130 - audio_seqlens: Optional[torch.Tensor] = None,
1131 - second_per_grids: Optional[torch.Tensor] = None,
1132 - ) -> tuple[torch.Tensor, torch.Tensor]:
1133 - def _get_feat_extract_output_lengths(input_lengths: torch.LongTensor):
1134 - input_lengths_leave = input_lengths % 100
1135 - feat_lengths = (input_lengths_leave - 1) // 2 + 1
1136 - output_lengths = (
1137 - ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
1138 - )
1139 - return output_lengths
1140 -
1141 - if input_ids is None or input_ids.ndim != 1:
1142 - raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids")
1143 -
1144 - seq_len = input_ids.shape[0]
1145 - device = input_ids.device
1146 - dtype = input_ids.dtype
1147 -
1148 - if image_grid_thw is not None:
1149 - image_grid_thw = image_grid_thw.to(device=device, dtype=torch.long)
1150 - if video_grid_thw is not None:
1151 - video_grid_thw = video_grid_thw.to(device=device, dtype=torch.long)
1152 -
1153 - if second_per_grids is None:
1154 - if video_grid_thw is not None and video_grid_thw.numel() > 0:
1155 - second_per_grids = torch.ones(
1156 - video_grid_thw.shape[0], dtype=torch.float32, device=device
1157 - )
1158 - else:
1159 - second_per_grids = torch.tensor([], dtype=torch.float32, device=device)
1160 - else:
1161 - second_per_grids = second_per_grids.to(device=device, dtype=torch.float32)
1162 -
1163 - if audio_seqlens is not None:
1164 - audio_seqlens = audio_seqlens.to(device=device, dtype=torch.long)
1165 -
1166 - spatial_merge_size = config.vision_config.spatial_merge_size
1167 - image_token_id = config.image_token_id
1168 - video_token_id = config.video_token_id
1169 - audio_token_id = config.audio_token_id
1170 - vision_start_token_id = config.vision_start_token_id
1171 - audio_start_token_id = config.audio_start_token_id
1172 - position_id_per_seconds = config.position_id_per_seconds
1173 -
1174 - vision_start_indices = torch.argwhere(
1175 - input_ids == vision_start_token_id
1176 - ).squeeze(1)
1177 - if vision_start_indices.numel() > 0:
1178 - vision_tokens = input_ids[vision_start_indices + 1]
1179 - else:
1180 - vision_tokens = input_ids.new_empty((0,), dtype=input_ids.dtype)
1181 - audio_nums = torch.sum(input_ids == audio_start_token_id)
1182 - image_nums = (vision_tokens == image_token_id).sum()
1183 - video_nums = (
1184 - (vision_tokens == audio_start_token_id).sum()
1185 - if use_audio_in_video
1186 - else (vision_tokens == video_token_id).sum()
1187 - )
1188 -
1189 - input_tokens = input_ids.tolist()
1190 - llm_pos_ids_list: list[torch.Tensor] = []
1191 - st = 0
1192 - image_idx = 0
1193 - video_idx = 0
1194 - audio_idx = 0
1195 - remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums # noqa: E501
1196 - multimodal_nums = (
1197 - image_nums + audio_nums
1198 - if use_audio_in_video
1199 - else image_nums + video_nums + audio_nums
1200 - ) # noqa: E501
1201 -
1202 - for _ in range(multimodal_nums):
1203 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1204 - if (image_token_id in input_tokens or video_token_id in input_tokens) and (
1205 - remain_videos > 0 or remain_images > 0
1206 - ):
1207 - ed_vision_start = input_tokens.index(vision_start_token_id, st)
1208 - else:
1209 - ed_vision_start = len(input_tokens) + 1
1210 - if audio_token_id in input_tokens and remain_audios > 0:
1211 - ed_audio_start = input_tokens.index(audio_start_token_id, st)
1212 - else:
1213 - ed_audio_start = len(input_tokens) + 1
1214 - min_ed = min(ed_vision_start, ed_audio_start)
1215 -
1216 - if min_ed == ed_audio_start:
1217 - text_len = min_ed - st
1218 - if text_len != 0:
1219 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1220 - llm_pos_ids_list.append(
1221 - torch.arange(text_len, device=device, dtype=torch.long)
1222 - .view(1, -1)
1223 - .expand(3, -1)
1224 - + st_idx
1225 - )
1226 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1227 - bos_len = 1
1228 - llm_pos_ids_list.append(
1229 - torch.arange(bos_len, device=device, dtype=torch.long)
1230 - .view(1, -1)
1231 - .expand(3, -1)
1232 - + st_idx
1233 - )
1234 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1235 - audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_idx])
1236 - llm_pos_ids = (
1237 - torch.arange(audio_len, device=device, dtype=torch.long)
1238 - .view(1, -1)
1239 - .expand(3, -1)
1240 - + st_idx
1241 - )
1242 - llm_pos_ids_list.append(llm_pos_ids)
1243 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1244 - eos_len = 1
1245 - llm_pos_ids_list.append(
1246 - torch.arange(eos_len, device=device, dtype=torch.long)
1247 - .view(1, -1)
1248 - .expand(3, -1)
1249 - + st_idx
1250 - )
1251 - st += text_len + bos_len + audio_len + eos_len
1252 - audio_idx += 1
1253 - remain_audios -= 1
1254 - elif (
1255 - min_ed == ed_vision_start
1256 - and input_ids[ed_vision_start + 1] == image_token_id
1257 - ):
1258 - text_len = min_ed - st
1259 - if text_len != 0:
1260 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1261 - llm_pos_ids_list.append(
1262 - torch.arange(text_len, device=device, dtype=torch.long)
1263 - .view(1, -1)
1264 - .expand(3, -1)
1265 - + st_idx
1266 - )
1267 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1268 - bos_len = 1
1269 - llm_pos_ids_list.append(
1270 - torch.arange(bos_len, device=device, dtype=torch.long)
1271 - .view(1, -1)
1272 - .expand(3, -1)
1273 - + st_idx
1274 - )
1275 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1276 - grid_t = image_grid_thw[image_idx][0]
1277 - grid_hs = image_grid_thw[:, 1]
1278 - grid_ws = image_grid_thw[:, 2]
1279 - t_index = torch.arange(grid_t, device=device) * position_id_per_seconds
1280 - llm_pos_ids = cls._get_llm_pos_ids_for_vision(
1281 - st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
1282 - )
1283 - image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
1284 - llm_pos_ids_list.append(llm_pos_ids)
1285 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1286 - eos_len = 1
1287 - llm_pos_ids_list.append(
1288 - torch.arange(eos_len, device=device, dtype=torch.long)
1289 - .view(1, -1)
1290 - .expand(3, -1)
1291 - + st_idx
1292 - )
1293 - st += text_len + bos_len + image_len + eos_len
1294 - image_idx += 1
1295 - remain_images -= 1
1296 - elif (
1297 - min_ed == ed_vision_start
1298 - and input_ids[ed_vision_start + 1] == video_token_id
1299 - and not use_audio_in_video
1300 - ):
1301 - text_len = min_ed - st
1302 - if text_len != 0:
1303 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1304 - llm_pos_ids_list.append(
1305 - torch.arange(text_len, device=device, dtype=torch.long)
1306 - .view(1, -1)
1307 - .expand(3, -1)
1308 - + st_idx
1309 - )
1310 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1311 - bos_len = 1
1312 - llm_pos_ids_list.append(
1313 - torch.arange(bos_len, device=device, dtype=torch.long)
1314 - .view(1, -1)
1315 - .expand(3, -1)
1316 - + st_idx
1317 - )
1318 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1319 - grid_t = video_grid_thw[video_idx][0]
1320 - grid_hs = video_grid_thw[:, 1]
1321 - grid_ws = video_grid_thw[:, 2]
1322 - t_index = (
1323 - torch.arange(grid_t, device=device)
1324 - * float(second_per_grids[video_idx].item())
1325 - * position_id_per_seconds
1326 - )
1327 - llm_pos_ids = cls._get_llm_pos_ids_for_vision(
1328 - st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
1329 - )
1330 - video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
1331 - llm_pos_ids_list.append(llm_pos_ids)
1332 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1333 - eos_len = 1
1334 - llm_pos_ids_list.append(
1335 - torch.arange(eos_len, device=device, dtype=torch.long)
1336 - .view(1, -1)
1337 - .expand(3, -1)
1338 - + st_idx
1339 - )
1340 - st += text_len + bos_len + video_len + eos_len
1341 - video_idx += 1
1342 - remain_videos -= 1
1343 - elif (
1344 - min_ed == ed_vision_start
1345 - and ed_vision_start + 1 == ed_audio_start
1346 - and use_audio_in_video
1347 - ):
1348 - text_len = min_ed - st
1349 - if text_len != 0:
1350 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1351 - llm_pos_ids_list.append(
1352 - torch.arange(text_len, device=device, dtype=torch.long)
1353 - .view(1, -1)
1354 - .expand(3, -1)
1355 - + st_idx
1356 - )
1357 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1358 - bos_len = 1
1359 - bos_block = (
1360 - torch.arange(bos_len, device=device, dtype=torch.long)
1361 - .view(1, -1)
1362 - .expand(3, -1)
1363 - + st_idx
1364 - )
1365 - llm_pos_ids_list.append(bos_block)
1366 - llm_pos_ids_list.append(bos_block)
1367 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1368 - audio_len = _get_feat_extract_output_lengths(audio_seqlens[audio_idx])
1369 - audio_llm_pos_ids = (
1370 - torch.arange(audio_len, device=device, dtype=torch.long)
1371 - .view(1, -1)
1372 - .expand(3, -1)
1373 - + st_idx
1374 - )
1375 - grid_t = video_grid_thw[video_idx][0]
1376 - grid_hs = video_grid_thw[:, 1]
1377 - grid_ws = video_grid_thw[:, 2]
1378 - t_index = (
1379 - torch.arange(grid_t, device=device)
1380 - * float(second_per_grids[video_idx].item())
1381 - * position_id_per_seconds
1382 - )
1383 - video_llm_pos_ids = cls._get_llm_pos_ids_for_vision(
1384 - st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
1385 - )
1386 - video_data_index, audio_data_index = 0, 0
1387 - while (
1388 - video_data_index < video_llm_pos_ids.shape[-1]
1389 - and audio_data_index < audio_llm_pos_ids.shape[-1]
1390 - ):
1391 - if (
1392 - video_llm_pos_ids[0][video_data_index]
1393 - <= audio_llm_pos_ids[0][audio_data_index]
1394 - ):
1395 - llm_pos_ids_list.append(
1396 - video_llm_pos_ids[
1397 - :, video_data_index : video_data_index + 1
1398 - ]
1399 - )
1400 - video_data_index += 1
1401 - else:
1402 - llm_pos_ids_list.append(
1403 - audio_llm_pos_ids[
1404 - :, audio_data_index : audio_data_index + 1
1405 - ]
1406 - )
1407 - audio_data_index += 1
1408 - if video_data_index < video_llm_pos_ids.shape[-1]:
1409 - llm_pos_ids_list.append(
1410 - video_llm_pos_ids[
1411 - :, video_data_index : video_llm_pos_ids.shape[-1]
1412 - ]
1413 - )
1414 - if audio_data_index < audio_llm_pos_ids.shape[-1]:
1415 - llm_pos_ids_list.append(
1416 - audio_llm_pos_ids[
1417 - :, audio_data_index : audio_llm_pos_ids.shape[-1]
1418 - ]
1419 - )
1420 - video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
1421 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1422 - eos_len = 1
1423 - eos_block = (
1424 - torch.arange(eos_len, device=device, dtype=torch.long)
1425 - .view(1, -1)
1426 - .expand(3, -1)
1427 - + st_idx
1428 - )
1429 - llm_pos_ids_list.append(eos_block)
1430 - llm_pos_ids_list.append(eos_block)
1431 - st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2 # noqa: E501
1432 - audio_idx += 1
1433 - video_idx += 1
1434 - remain_videos -= 1
1435 - remain_audios -= 1
1436 -
1437 - if st < len(input_tokens):
1438 - st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1439 - text_len = len(input_tokens) - st
1440 - llm_pos_ids_list.append(
1441 - torch.arange(text_len, device=device, dtype=torch.long)
1442 - .view(1, -1)
1443 - .expand(3, -1)
1444 - + st_idx
1445 - )
1446 -
1447 - llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1448 - if llm_positions.shape[1] != seq_len:
1449 - raise RuntimeError("Position ids length mismatch with input ids length")
1450 -
1451 - position_ids = llm_positions.to(device=device, dtype=dtype)
1452 - mrope_position_delta = llm_positions.max() + 1 - seq_len
1453 - return position_ids, mrope_position_delta
1454 -
1122 @classmethod
1123 def _omni_get_input_positions_tensor(
1124 cls,
@@ -1483,8 +1150,6 @@ class MRotaryEmbedding(RotaryEmbedding):
1150
1151 # TODO(fyabc): refactor and share more code with
1152 # _vl_get_input_positions_tensor.
1486 -
1487 - model_type = hf_config.model_type
1153 thinker_config = hf_config.thinker_config
1154
1155 if isinstance(image_grid_thw, list):
@@ -1492,30 +1157,6 @@ class MRotaryEmbedding(RotaryEmbedding):
1157 if isinstance(video_grid_thw, list):
1158 video_grid_thw = torch.tensor(video_grid_thw)
1159
1495 - if "qwen3_omni" in model_type:
1496 - input_tensor = torch.tensor(input_tokens)
1497 - audio_lengths_tensor = audio_feature_lengths
1498 - if audio_lengths_tensor is not None and not isinstance(
1499 - audio_lengths_tensor, torch.Tensor
1500 - ):
1501 - audio_lengths_tensor = torch.as_tensor(
1502 - audio_lengths_tensor, dtype=torch.long
1503 - )
1504 - second_per_grids_tensor = (
1505 - torch.tensor(second_per_grid_ts) if second_per_grid_ts else None
1506 - )
1507 -
1508 - llm_positions, mrope_position_delta = cls._omni3_get_input_positions_tensor( # noqa: E501
1509 - thinker_config,
1510 - input_tensor,
1511 - image_grid_thw,
1512 - video_grid_thw,
1513 - use_audio_in_video,
1514 - audio_lengths_tensor,
1515 - second_per_grids_tensor,
1516 - )
1517 - return llm_positions, mrope_position_delta
1518 -
1160 audio_token_id = thinker_config.audio_token_index
1161 image_token_id = thinker_config.image_token_index
1162 video_token_id = thinker_config.video_token_index
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+329 -26
@@ -72,7 +72,12 @@ from vllm.multimodal.processing import (
72 )
73 from vllm.sequence import IntermediateTensors
74
75 -from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP
75 +from .interfaces import (
76 + MultiModalEmbeddings,
77 + SupportsMRoPE,
78 + SupportsMultiModal,
79 + SupportsPP,
80 +)
81
82 # yapf conflicts with isort for this block
83 # yapf: disable
@@ -96,7 +101,7 @@ from .utils import (
101 _merge_multimodal_embeddings,
102 maybe_prefix,
103 )
99 -from .vision import get_vit_attn_backend
104 +from .vision import get_llm_pos_ids_for_vision, get_vit_attn_backend
105
106 try:
107 import flash_attn
@@ -106,6 +111,15 @@ except (ImportError, ModuleNotFoundError):
111 logger = init_logger(__name__)
112
113
114 +def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
115 + input_lengths_leave = input_lengths % 100
116 + feat_lengths = (input_lengths_leave - 1) // 2 + 1
117 + output_lengths = (
118 + ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
119 + )
120 + return feat_lengths, output_lengths
121 +
122 +
123 class Qwen3_VisionPatchEmbed(nn.Module):
124 def __init__(
125 self,
@@ -679,16 +693,6 @@ Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder
693 class Qwen3OmniMoeThinkerMultiModalProcessor(
694 Qwen2_5OmniThinkerMultiModalProcessor,
695 ):
682 - def _get_feat_extract_output_lengths(
683 - self, input_lengths: torch.Tensor
684 - ) -> torch.Tensor:
685 - input_lengths_leave = input_lengths % 100
686 - feat_lengths = (input_lengths_leave - 1) // 2 + 1
687 - output_lengths = (
688 - ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
689 - )
690 - return feat_lengths, output_lengths
691 -
696 def _call_hf_processor(
697 self,
698 prompt: str,
@@ -882,13 +886,13 @@ class Qwen3OmniMoeThinkerMultiModalProcessor(
886 if audio_feature_lengths is None and feature_attention_mask is None:
887 audio_output_lengths = []
888 elif audio_feature_lengths is not None:
885 - _, audio_output_lens = self._get_feat_extract_output_lengths(
889 + _, audio_output_lens = _get_feat_extract_output_lengths(
890 audio_feature_lengths
891 )
892 audio_output_lengths = audio_output_lens.tolist()
893 elif feature_attention_mask is not None:
894 assert isinstance(feature_attention_mask, torch.Tensor)
891 - _, audio_output_lens = self._get_feat_extract_output_lengths(
895 + _, audio_output_lens = _get_feat_extract_output_lengths(
896 feature_attention_mask.sum(-1)
897 )
898 audio_output_lengths = audio_output_lens.tolist()
@@ -1044,16 +1048,6 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
1048 else:
1049 return torch.concat(mm_input, dim=dim)
1050
1047 - def _get_feat_extract_output_lengths(
1048 - self, input_lengths: torch.Tensor
1049 - ) -> torch.Tensor:
1050 - input_lengths_leave = input_lengths % 100
1051 - feat_lengths = (input_lengths_leave - 1) // 2 + 1
1052 - output_lengths = (
1053 - ((feat_lengths - 1) // 2 + 1 - 1) // 2 + 1 + (input_lengths // 100) * 13
1054 - )
1055 - return output_lengths, output_lengths
1056 -
1051 def _process_audio_input(
1052 self,
1053 audio_input: Qwen2AudioFeatureInputs,
@@ -1072,8 +1066,8 @@ class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMix
1066 if audio_feature_lengths.ndim == 2:
1067 audio_feature_lengths = audio_feature_lengths.reshape(-1)
1068
1075 - audio_feat_lengths, audio_output_lengths = (
1076 - self._get_feat_extract_output_lengths(audio_feature_lengths)
1069 + audio_feat_lengths, audio_output_lengths = _get_feat_extract_output_lengths(
1070 + audio_feature_lengths
1071 )
1072
1073 audio_outputs = self.audio_tower(
@@ -1094,6 +1088,7 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1088 nn.Module,
1089 SupportsMultiModal,
1090 SupportsPP,
1091 + SupportsMRoPE,
1092 Qwen3OmniMoeConditionalGenerationMixin,
1093 ):
1094 hf_to_vllm_mapper = WeightsMapper(
@@ -1407,3 +1402,311 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1402 loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)
1403
1404 return loaded_weights
1405 +
1406 + @classmethod
1407 + def get_mrope_input_positions(
1408 + self,
1409 + input_tokens: list[int],
1410 + hf_config: PretrainedConfig,
1411 + image_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
1412 + video_grid_thw: Optional[Union[list[list[int]], torch.Tensor]],
1413 + second_per_grid_ts: Optional[list[float]] = None,
1414 + context_len: int = 0,
1415 + seq_len: Optional[int] = None,
1416 + audio_feature_lengths: Optional[torch.Tensor] = None,
1417 + use_audio_in_video: bool = False,
1418 + ) -> tuple[torch.Tensor, int]:
1419 + config = hf_config.thinker_config
1420 + if isinstance(image_grid_thw, list):
1421 + image_grid_thw = torch.tensor(image_grid_thw)
1422 + if isinstance(video_grid_thw, list):
1423 + video_grid_thw = torch.tensor(video_grid_thw)
1424 + input_ids = torch.tensor(input_tokens)
1425 + if input_ids is None or input_ids.ndim != 1:
1426 + raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids")
1427 +
1428 + seq_len = input_ids.shape[0]
1429 + if audio_feature_lengths is not None and not isinstance(
1430 + audio_feature_lengths, torch.Tensor
1431 + ):
1432 + audio_feature_lengths = torch.as_tensor(
1433 + audio_feature_lengths, dtype=torch.long
1434 + )
1435 + if second_per_grid_ts is None:
1436 + if video_grid_thw is not None and video_grid_thw.numel() > 0:
1437 + second_per_grids = torch.ones(
1438 + video_grid_thw.shape[0], dtype=torch.float32
1439 + )
1440 + else:
1441 + second_per_grids = torch.tensor([], dtype=torch.float32)
1442 + else:
1443 + second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32)
1444 +
1445 + spatial_merge_size = config.vision_config.spatial_merge_size
1446 + image_token_id = config.image_token_id
1447 + video_token_id = config.video_token_id
1448 + audio_token_id = config.audio_token_id
1449 + vision_start_token_id = config.vision_start_token_id
1450 + audio_start_token_id = config.audio_start_token_id
1451 + position_id_per_seconds = config.position_id_per_seconds
1452 +
1453 + vision_start_indices = torch.argwhere(
1454 + input_ids == vision_start_token_id
1455 + ).squeeze(1)
1456 + if vision_start_indices.numel() > 0:
1457 + vision_tokens = input_ids[vision_start_indices + 1]
1458 + else:
1459 + vision_tokens = input_ids.new_empty((0,), dtype=input_ids.dtype)
1460 + audio_nums = torch.sum(input_ids == audio_start_token_id)
1461 + image_nums = (vision_tokens == image_token_id).sum()
1462 + video_nums = (
1463 + (vision_tokens == audio_start_token_id).sum()
1464 + if use_audio_in_video
1465 + else (vision_tokens == video_token_id).sum()
1466 + )
1467 +
1468 + llm_pos_ids_list: list[torch.Tensor] = []
1469 + st = 0
1470 + image_idx = 0
1471 + video_idx = 0
1472 + audio_idx = 0
1473 + remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums # noqa: E501
1474 + multimodal_nums = (
1475 + image_nums + audio_nums
1476 + if use_audio_in_video
1477 + else image_nums + video_nums + audio_nums
1478 + ) # noqa: E501
1479 +
1480 + for _ in range(multimodal_nums):
1481 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1482 + if (image_token_id in input_tokens or video_token_id in input_tokens) and (
1483 + remain_videos > 0 or remain_images > 0
1484 + ):
1485 + ed_vision_start = input_tokens.index(vision_start_token_id, st)
1486 + else:
1487 + ed_vision_start = len(input_tokens) + 1
1488 + if audio_token_id in input_tokens and remain_audios > 0:
1489 + ed_audio_start = input_tokens.index(audio_start_token_id, st)
1490 + else:
1491 + ed_audio_start = len(input_tokens) + 1
1492 + min_ed = min(ed_vision_start, ed_audio_start)
1493 +
1494 + if min_ed == ed_audio_start:
1495 + text_len = min_ed - st
1496 + if text_len != 0:
1497 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1498 + llm_pos_ids_list.append(
1499 + torch.arange(text_len, dtype=torch.long)
1500 + .view(1, -1)
1501 + .expand(3, -1)
1502 + + st_idx
1503 + )
1504 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1505 + bos_len = 1
1506 + llm_pos_ids_list.append(
1507 + torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1508 + + st_idx
1509 + )
1510 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1511 + _, audio_len = _get_feat_extract_output_lengths(
1512 + audio_feature_lengths[audio_idx]
1513 + )
1514 + llm_pos_ids = (
1515 + torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1)
1516 + + st_idx
1517 + )
1518 + llm_pos_ids_list.append(llm_pos_ids)
1519 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1520 + eos_len = 1
1521 + llm_pos_ids_list.append(
1522 + torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1523 + + st_idx
1524 + )
1525 + st += text_len + bos_len + audio_len + eos_len
1526 + audio_idx += 1
1527 + remain_audios -= 1
1528 + elif (
1529 + min_ed == ed_vision_start
1530 + and input_ids[ed_vision_start + 1] == image_token_id
1531 + ):
1532 + text_len = min_ed - st
1533 + if text_len != 0:
1534 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1535 + llm_pos_ids_list.append(
1536 + torch.arange(text_len, dtype=torch.long)
1537 + .view(1, -1)
1538 + .expand(3, -1)
1539 + + st_idx
1540 + )
1541 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1542 + bos_len = 1
1543 + llm_pos_ids_list.append(
1544 + torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1545 + + st_idx
1546 + )
1547 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1548 + grid_t = image_grid_thw[image_idx][0]
1549 + grid_hs = image_grid_thw[:, 1]
1550 + grid_ws = image_grid_thw[:, 2]
1551 + t_index = torch.arange(grid_t) * position_id_per_seconds
1552 + llm_pos_ids = get_llm_pos_ids_for_vision(
1553 + st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
1554 + )
1555 + image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
1556 + llm_pos_ids_list.append(llm_pos_ids)
1557 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1558 + eos_len = 1
1559 + llm_pos_ids_list.append(
1560 + torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1561 + + st_idx
1562 + )
1563 + st += text_len + bos_len + image_len + eos_len
1564 + image_idx += 1
1565 + remain_images -= 1
1566 + elif (
1567 + min_ed == ed_vision_start
1568 + and input_ids[ed_vision_start + 1] == video_token_id
1569 + and not use_audio_in_video
1570 + ):
1571 + text_len = min_ed - st
1572 + if text_len != 0:
1573 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1574 + llm_pos_ids_list.append(
1575 + torch.arange(text_len, dtype=torch.long)
1576 + .view(1, -1)
1577 + .expand(3, -1)
1578 + + st_idx
1579 + )
1580 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1581 + bos_len = 1
1582 + llm_pos_ids_list.append(
1583 + torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1584 + + st_idx
1585 + )
1586 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1587 + grid_t = video_grid_thw[video_idx][0]
1588 + grid_hs = video_grid_thw[:, 1]
1589 + grid_ws = video_grid_thw[:, 2]
1590 + t_index = (
1591 + torch.arange(grid_t)
1592 + * float(second_per_grids[video_idx].item())
1593 + * position_id_per_seconds
1594 + )
1595 + llm_pos_ids = get_llm_pos_ids_for_vision(
1596 + st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
1597 + )
1598 + video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
1599 + llm_pos_ids_list.append(llm_pos_ids)
1600 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1601 + eos_len = 1
1602 + llm_pos_ids_list.append(
1603 + torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1604 + + st_idx
1605 + )
1606 + st += text_len + bos_len + video_len + eos_len
1607 + video_idx += 1
1608 + remain_videos -= 1
1609 + elif (
1610 + min_ed == ed_vision_start
1611 + and ed_vision_start + 1 == ed_audio_start
1612 + and use_audio_in_video
1613 + ):
1614 + text_len = min_ed - st
1615 + if text_len != 0:
1616 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1617 + llm_pos_ids_list.append(
1618 + torch.arange(text_len, dtype=torch.long)
1619 + .view(1, -1)
1620 + .expand(3, -1)
1621 + + st_idx
1622 + )
1623 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1624 + bos_len = 1
1625 + bos_block = (
1626 + torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1627 + + st_idx
1628 + )
1629 + llm_pos_ids_list.append(bos_block)
1630 + llm_pos_ids_list.append(bos_block)
1631 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1632 + _, audio_len = _get_feat_extract_output_lengths(
1633 + audio_feature_lengths[audio_idx]
1634 + )
1635 + audio_llm_pos_ids = (
1636 + torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1)
1637 + + st_idx
1638 + )
1639 + grid_t = video_grid_thw[video_idx][0]
1640 + grid_hs = video_grid_thw[:, 1]
1641 + grid_ws = video_grid_thw[:, 2]
1642 + t_index = (
1643 + torch.arange(grid_t)
1644 + * float(second_per_grids[video_idx].item())
1645 + * position_id_per_seconds
1646 + )
1647 + video_llm_pos_ids = get_llm_pos_ids_for_vision(
1648 + st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
1649 + )
1650 + video_data_index, audio_data_index = 0, 0
1651 + while (
1652 + video_data_index < video_llm_pos_ids.shape[-1]
1653 + and audio_data_index < audio_llm_pos_ids.shape[-1]
1654 + ):
1655 + if (
1656 + video_llm_pos_ids[0][video_data_index]
1657 + <= audio_llm_pos_ids[0][audio_data_index]
1658 + ):
1659 + llm_pos_ids_list.append(
1660 + video_llm_pos_ids[
1661 + :, video_data_index : video_data_index + 1
1662 + ]
1663 + )
1664 + video_data_index += 1
1665 + else:
1666 + llm_pos_ids_list.append(
1667 + audio_llm_pos_ids[
1668 + :, audio_data_index : audio_data_index + 1
1669 + ]
1670 + )
1671 + audio_data_index += 1
1672 + if video_data_index < video_llm_pos_ids.shape[-1]:
1673 + llm_pos_ids_list.append(
1674 + video_llm_pos_ids[
1675 + :, video_data_index : video_llm_pos_ids.shape[-1]
1676 + ]
1677 + )
1678 + if audio_data_index < audio_llm_pos_ids.shape[-1]:
1679 + llm_pos_ids_list.append(
1680 + audio_llm_pos_ids[
1681 + :, audio_data_index : audio_llm_pos_ids.shape[-1]
1682 + ]
1683 + )
1684 + video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
1685 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1686 + eos_len = 1
1687 + eos_block = (
1688 + torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1689 + + st_idx
1690 + )
1691 + llm_pos_ids_list.append(eos_block)
1692 + llm_pos_ids_list.append(eos_block)
1693 + st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2 # noqa: E501
1694 + audio_idx += 1
1695 + video_idx += 1
1696 + remain_videos -= 1
1697 + remain_audios -= 1
1698 +
1699 + if st < len(input_tokens):
1700 + st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1701 + text_len = len(input_tokens) - st
1702 + llm_pos_ids_list.append(
1703 + torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1)
1704 + + st_idx
1705 + )
1706 +
1707 + llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
1708 + if llm_positions.shape[1] != seq_len:
1709 + raise RuntimeError("Position ids length mismatch with input ids length")
1710 +
1711 + mrope_position_delta = llm_positions.max() + 1 - seq_len
1712 + return llm_positions, mrope_position_delta
vllm/model_executor/models/vision.py
+37
@@ -499,3 +499,40 @@ def run_dp_sharded_mrope_vision_model(
499 "Found unassigned embeddings"
500 )
501 return out_embeddings
502 +
503 +
504 +def get_llm_pos_ids_for_vision(
505 + start_idx: int,
506 + vision_idx: int,
507 + spatial_merge_size: int,
508 + t_index: list[int],
509 + grid_hs: torch.Tensor,
510 + grid_ws: torch.Tensor,
511 +) -> torch.Tensor:
512 + llm_pos_ids_list = []
513 + llm_grid_h = grid_hs[vision_idx] // spatial_merge_size
514 + llm_grid_w = grid_ws[vision_idx] // spatial_merge_size
515 + h_index = (
516 + torch.arange(llm_grid_h)
517 + .view(1, -1, 1)
518 + .expand(len(t_index), -1, llm_grid_w)
519 + .flatten()
520 + )
521 + w_index = (
522 + torch.arange(llm_grid_w)
523 + .view(1, 1, -1)
524 + .expand(len(t_index), llm_grid_h, -1)
525 + .flatten()
526 + )
527 + t_index_tensor = (
528 + torch.Tensor(t_index)
529 + .to(llm_grid_h.device)
530 + .view(-1, 1)
531 + .expand(-1, llm_grid_h * llm_grid_w)
532 + .long()
533 + .flatten()
534 + )
535 + _llm_pos_ids = torch.stack([t_index_tensor, h_index, w_index])
536 + llm_pos_ids_list.append(_llm_pos_ids + start_idx)
537 + llm_pos_ids = torch.cat(llm_pos_ids_list, dim=1)
538 + return llm_pos_ids
vllm/v1/worker/gpu_model_runner.py
+1 -1
@@ -875,7 +875,7 @@ class GPUModelRunner(LoRAModelRunnerMixin, KVConnectorModelRunnerMixin):
875 if mm_input.get("use_audio_in_video") is True:
876 use_audio_in_video = True
877
878 - if supports_mrope(self.model):
878 + if supports_mrope(self.get_model()):
879 req_state.mrope_positions, req_state.mrope_position_delta = (
880 self.model.get_mrope_input_positions(
881 req_state.prompt_token_ids,