@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 19a9b169bf

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

Add Qwen3-Omni moe thinker (#25550)

Signed-off-by: DarkLight1337 <tlleungac@connect.ust.hk> Signed-off-by: Roger Wang <hey@rogerw.io> Signed-off-by: Isotr0py <mozf@mail2.sysu.edu.cn> Co-authored-by: Xiong Wang <feizi.wx@alibaba-inc.com> Co-authored-by: DarkLight1337 <tlleungac@connect.ust.hk> Co-authored-by: Roger Wang <hey@rogerw.io> Co-authored-by: Isotr0py <mozf@mail2.sysu.edu.cn>

Xiong Wang committed Oct 11, 2025 at 01:00 UTC 19a9b169bf1b58b4311f5f3bfa37550328741506

6 files changed +1795 -36

docs/models/supported_models.md

+2 -2

@@ -714,6 +714,7 @@ These models primarily accept the [`LLM.generate`](./generative_models.md#llmgen

714 | `Qwen2_5OmniThinkerForConditionalGeneration` | Qwen2.5-Omni | T + IE+ + VE+ + A+ | `Qwen/Qwen2.5-Omni-3B`, `Qwen/Qwen2.5-Omni-7B` | ✅︎ | ✅︎ | ✅︎ |

715 | `Qwen3VLForConditionalGeneration` | Qwen3-VL | T + IE+ + VE+ | `Qwen/Qwen3-VL-4B-Instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |

716 | `Qwen3VLMoeForConditionalGeneration` | Qwen3-VL-MOE | T + IE+ + VE+ | `Qwen/Qwen3-VL-30B-A3B-Instruct`, etc. | ✅︎ | ✅︎ | ✅︎ |

717 +| `Qwen3OmniMoeThinkerForConditionalGeneration` | Qwen3-Omni | T + IE+ + VE+ + A+ | `Qwen/Qwen3-Omni-30B-A3B-Instruct`, `Qwen/Qwen3-Omni-30B-A3B-Thinking` | ✅︎ | ✅︎ | ✅︎ |

718 | `RForConditionalGeneration` | R-VL-4B | T + IE+ | `YannQi/R-4B` | | ✅︎ | ✅︎ |

719 | `SkyworkR1VChatModel` | Skywork-R1V-38B | T + I | `Skywork/Skywork-R1V-38B` | | ✅︎ | ✅︎ |

720 | `SmolVLMForConditionalGeneration` | SmolVLM2 | T + I | `SmolVLM2-2.2B-Instruct` | ✅︎ | | ✅︎ |

@@ -804,8 +805,7 @@ Some models are supported only via the [Transformers backend](#transformers). Th

805 Our PaliGemma implementations have the same problem as Gemma 3 (see above) for both V0 and V1.

806

807 !!! note

807 - For Qwen2.5-Omni, reading audio from video pre-processing (`--mm-processor-kwargs '{"use_audio_in_video": true}'`)

808 - is currently supported on V0 (but not V1), because overlapping modalities is not yet supported in V1.

808 + For Qwen2.5-Omni and Qwen3-Omni, reading audio from video pre-processing (`--mm-processor-kwargs '{"use_audio_in_video": true}'`) is currently work in progress and not yet supported.

809

810 #### Transcription

811

tests/models/multimodal/processing/test_common.py

@@ -384,6 +384,7 @@ def _test_processing_correctness_one(

384 "Qwen/Qwen2.5-Omni-3B",

385 "Qwen/Qwen3-VL-4B-Instruct",

386 "Qwen/Qwen3-VL-30B-A3B-Instruct",

387 + "Qwen/Qwen3-Omni-30B-A3B-Instruct",

388 "YannQi/R-4B",

389 "Skywork/Skywork-R1V-38B",

390 "HuggingFaceTB/SmolVLM2-2.2B-Instruct",

tests/models/registry.py

@@ -773,6 +773,11 @@ _MULTIMODAL_EXAMPLE_MODELS = {

773 min_transformers_version="4.57",

774 is_available_online=False,

775 ),

776 + "Qwen3OmniMoeForConditionalGeneration": _HfExamplesInfo(

777 + "Qwen/Qwen3-Omni-30B-A3B-Instruct",

778 + max_model_len=4096,

779 + min_transformers_version="4.57",

780 + ),

781 "RForConditionalGeneration": _HfExamplesInfo("YannQi/R-4B", trust_remote_code=True),

782 "SkyworkR1VChatModel": _HfExamplesInfo(

783 "Skywork/Skywork-R1V-38B", trust_remote_code=True

vllm/model_executor/layers/rotary_embedding/mrope.py

+374 -34

@@ -971,17 +971,9 @@ class MRotaryEmbedding(RotaryEmbedding):

971 llm_pos_ids_list.append(

972 torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx

973 )

974 -

974 t_index = (

976 - (

977 - torch.arange(llm_grid_t)

978 - .view(-1, 1)

979 - .expand(-1, llm_grid_h * llm_grid_w)

980 - )

981 - .long()

982 - .flatten()

983 - )

984 -

975 + torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)

976 + ).flatten()

977 h_index = (

978 torch.arange(llm_grid_h)

979 .view(1, -1, 1)

@@ -1042,7 +1034,6 @@ class MRotaryEmbedding(RotaryEmbedding):

1034

1035 st = 0

1036 remain_images, remain_videos = image_nums, video_nums

1045 -

1037 image_index, video_index = 0, 0

1038 for _ in range(image_nums + video_nums):

1039 video_second_per_grid_t = 0.0

@@ -1093,19 +1084,11 @@ class MRotaryEmbedding(RotaryEmbedding):

1084 llm_pos_ids_list.append(

1085 torch.arange(text_len).view(1, -1).expand(3, -1) + st_idx

1086 )

1096 -

1087 t_index = (

1098 - (

1099 - torch.arange(llm_grid_t)

1100 - .view(-1, 1)

1101 - .expand(-1, llm_grid_h * llm_grid_w)

1102 - * video_second_per_grid_t

1103 - * tokens_per_second

1104 - )

1105 - .long()

1106 - .flatten()

1107 - )

1108 -

1088 + torch.arange(llm_grid_t).view(-1, 1).expand(-1, llm_grid_h * llm_grid_w)

1089 + * video_second_per_grid_t

1090 + * tokens_per_second

1091 + ).flatten()

1092 h_index = (

1093 torch.arange(llm_grid_h)

1094 .view(1, -1, 1)

@@ -1136,6 +1119,339 @@ 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 +

1455 @classmethod

1456 def _omni_get_input_positions_tensor(

1457 cls,

@@ -1168,7 +1484,38 @@ class MRotaryEmbedding(RotaryEmbedding):

1484 # TODO(fyabc): refactor and share more code with

1485 # _vl_get_input_positions_tensor.

1486

1487 + model_type = hf_config.model_type

1488 thinker_config = hf_config.thinker_config

1489 +

1490 + if isinstance(image_grid_thw, list):

1491 + image_grid_thw = torch.tensor(image_grid_thw)

1492 + if isinstance(video_grid_thw, list):

1493 + video_grid_thw = torch.tensor(video_grid_thw)

1494 +

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 +

1519 audio_token_id = thinker_config.audio_token_index

1520 image_token_id = thinker_config.image_token_index

1521 video_token_id = thinker_config.video_token_index

@@ -1182,11 +1529,6 @@ class MRotaryEmbedding(RotaryEmbedding):

1529 thinker_config.vision_config, "tokens_per_second", 25

1530 )

1531

1185 - if isinstance(image_grid_thw, list):

1186 - image_grid_thw = torch.tensor(image_grid_thw)

1187 - if isinstance(video_grid_thw, list):

1188 - video_grid_thw = torch.tensor(video_grid_thw)

1189 -

1532 src_item = input_tokens

1533 audio_seqlens = audio_feature_lengths

1534 if not second_per_grid_ts:

@@ -1232,7 +1574,7 @@ class MRotaryEmbedding(RotaryEmbedding):

1574 grid_t = image_grid_thw[image_idx][0]

1575 grid_hs = image_grid_thw[:, 1]

1576 grid_ws = image_grid_thw[:, 2]

1235 - t_index = (torch.arange(grid_t) * 1 * tokens_per_second).long()

1577 + t_index = torch.arange(grid_t) * 1 * tokens_per_second

1578 llm_pos_ids = cls._get_llm_pos_ids_for_vision(

1579 start_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws

1580 )

@@ -1250,7 +1592,7 @@ class MRotaryEmbedding(RotaryEmbedding):

1592 torch.arange(grid_t)

1593 * second_per_grid_ts[video_idx]

1594 * tokens_per_second

1253 - ).long()

1595 + )

1596 llm_pos_ids = cls._get_llm_pos_ids_for_vision(

1597 start_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws

1598 )

@@ -1277,7 +1619,7 @@ class MRotaryEmbedding(RotaryEmbedding):

1619 torch.arange(grid_t)

1620 * second_per_grid_ts[video_idx]

1621 * tokens_per_second

1280 - ).long()

1622 + )

1623 t_index_split_chunk = cls._split_list_into_ranges(

1624 t_index, t_ntoken_per_chunk

1625 )

@@ -1452,9 +1794,7 @@ class MRotaryEmbedding(RotaryEmbedding):

1794 grid_h = video_grid_thw[1]

1795 grid_w = video_grid_thw[2]

1796 t_ntoken_per_chunk = int(tokens_per_second * seconds_per_chunk)

1455 - t_index = (

1456 - torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second

1457 - ).long()

1797 + t_index = torch.arange(grid_t) * video_second_per_grid_t * tokens_per_second

1798 t_index_split_chunk = cls._split_list_into_ranges(t_index, t_ntoken_per_chunk)

1799

1800 updates = [audio_start_token_id]

vllm/model_executor/models/qwen3_omni_moe_thinker.py new

+1409

@@ -0,0 +1,1409 @@

1 +# SPDX-License-Identifier: Apache-2.0

2 +# SPDX-FileCopyrightText: Copyright contributors to the vLLM project

6 +#

7 +# This code is based on EleutherAI's GPT-NeoX library and the GPT-NeoX

8 +# and OPT implementations in this library. It has been modified from its

9 +# original forms to accommodate minor architectural differences compared

10 +# to GPT-NeoX and OPT used by the Meta AI team that trained the model.

11 +#

12 +# Licensed under the Apache License, Version 2.0 (the "License");

13 +# you may not use this file except in compliance with the License.

14 +# You may obtain a copy of the License at

15 +#

16 +# http://www.apache.org/licenses/LICENSE-2.0

17 +#

18 +# Unless required by applicable law or agreed to in writing, software

19 +# distributed under the License is distributed on an "AS IS" BASIS,

20 +# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.

21 +# See the License for the specific language governing permissions and

22 +# limitations under the License.

23 +"""Inference-only Qwen3-Omni-Moe model (thinker part)."""

24 +

25 +from collections.abc import Iterable, Mapping, Sequence

26 +from functools import partial

27 +from typing import Any, Callable, Optional, Union

28 +

29 +import numpy as np

30 +import torch

31 +import torch.nn as nn

32 +import torch.nn.functional as F

33 +from transformers import PretrainedConfig

34 +from transformers.feature_extraction_utils import BatchFeature

35 +from transformers.models.qwen3_omni_moe.configuration_qwen3_omni_moe import (

36 + Qwen3OmniMoeConfig,

37 + Qwen3OmniMoeThinkerConfig,

38 +)

39 +from transformers.models.qwen3_omni_moe.modeling_qwen3_omni_moe import (

40 + Qwen3OmniMoeAudioEncoder,

41 +)

42 +from transformers.models.qwen3_omni_moe.processing_qwen3_omni_moe import (

43 + Qwen3OmniMoeProcessor,

44 +)

45 +from transformers.models.whisper import WhisperFeatureExtractor

46 +

47 +from vllm.attention.backends.registry import _Backend

48 +from vllm.attention.layer import check_upstream_fa_availability

49 +from vllm.compilation.decorators import support_torch_compile

50 +from vllm.config import VllmConfig

51 +from vllm.distributed import get_pp_group

52 +from vllm.logger import init_logger

53 +from vllm.model_executor.layers.activation import _ACTIVATION_REGISTRY

54 +from vllm.model_executor.layers.linear import ColumnParallelLinear, RowParallelLinear

55 +from vllm.model_executor.layers.logits_processor import LogitsProcessor

56 +from vllm.model_executor.layers.quantization import QuantizationConfig

57 +from vllm.model_executor.layers.vocab_parallel_embedding import ParallelLMHead

58 +from vllm.model_executor.model_loader.weight_utils import default_weight_loader

59 +from vllm.model_executor.models.qwen2_audio import (

60 + Qwen2AudioFeatureInputs,

61 + Qwen2AudioProcessingInfo,

62 +)

63 +from vllm.multimodal import MULTIMODAL_REGISTRY

64 +from vllm.multimodal.inputs import MultiModalKwargsItems

65 +from vllm.multimodal.parse import AudioProcessorItems, MultiModalDataItems

66 +from vllm.multimodal.processing import (

67 + BaseMultiModalProcessor,

68 + MultiModalPromptUpdates,

69 + PlaceholderFeaturesInfo,

70 + PromptReplacement,

71 + PromptUpdate,

72 +)

73 +from vllm.sequence import IntermediateTensors

74 +

75 +from .interfaces import MultiModalEmbeddings, SupportsMultiModal, SupportsPP

76 +

77 +# yapf conflicts with isort for this block

78 +# yapf: disable

79 +from .qwen2_5_omni_thinker import (

80 + Qwen2_5OmniConditionalGenerationMixin,

81 + Qwen2_5OmniThinkerDummyInputsBuilder,

82 + Qwen2_5OmniThinkerMultiModalProcessor,

83 + Qwen2_5OmniThinkerProcessingInfo,

84 +)

85 +

86 +# yapf: enable

87 +from .qwen2_5_vl import (

88 + Qwen2_5_VisionAttention,

89 + Qwen2_5_VisionRotaryEmbedding,

90 + Qwen2_5_VLProcessingInfo,

91 +)

92 +from .qwen3_moe import Qwen3MoeForCausalLM, Qwen3MoeModel

93 +from .utils import (

94 + AutoWeightsLoader,

95 + WeightsMapper,

96 + _merge_multimodal_embeddings,

97 + maybe_prefix,

98 +)

99 +from .vision import get_vit_attn_backend

100 +

101 +try:

102 + import flash_attn

103 +except (ImportError, ModuleNotFoundError):

104 + flash_attn = None

105 +

106 +logger = init_logger(__name__)

107 +

108 +

109 +class Qwen3_VisionPatchEmbed(nn.Module):

110 + def __init__(

111 + self,

112 + patch_size: int = 14,

113 + temporal_patch_size: int = 2,

114 + in_channels: int = 3,

115 + hidden_size: int = 1152,

116 + ) -> None:

117 + super().__init__()

118 + self.patch_size = patch_size

119 + self.temporal_patch_size = temporal_patch_size

120 + self.hidden_size = hidden_size

121 +

122 + kernel_size = (temporal_patch_size, patch_size, patch_size)

123 + self.proj = nn.Conv3d(

124 + in_channels,

125 + hidden_size,

126 + kernel_size=kernel_size,

127 + stride=kernel_size,

128 + bias=True,

129 + )

130 +

131 + def forward(self, x: torch.Tensor) -> torch.Tensor:

132 + L, C = x.shape

133 + x = x.view(L, -1, self.temporal_patch_size, self.patch_size, self.patch_size)

134 + x = self.proj(x).view(L, self.hidden_size)

135 + return x

136 +

137 +

138 +class Qwen3_VisionMLP(nn.Module):

139 + def __init__(

140 + self,

141 + in_features: int,

142 + hidden_features: int,

143 + bias: bool = False,

144 + act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

145 + quant_config: Optional[QuantizationConfig] = None,

146 + prefix: str = "",

147 + ):

148 + super().__init__()

149 + self.linear_fc1 = ColumnParallelLinear(

150 + in_features,

151 + hidden_features,

152 + bias=bias,

153 + quant_config=quant_config,

154 + return_bias=False,

155 + prefix=f"{prefix}.linear_fc1",

156 + )

157 + self.linear_fc2 = RowParallelLinear(

158 + hidden_features,

159 + in_features,

160 + bias=bias,

161 + quant_config=quant_config,

162 + return_bias=False,

163 + prefix=f"{prefix}.linear_fc2",

164 + )

165 + self.act_fn = act_fn

166 +

167 + def forward(self, x: torch.Tensor):

168 + mlp_output = self.linear_fc2(self.act_fn(self.linear_fc1(x)))

169 + return mlp_output

170 +

171 +

172 +class Qwen3_VisionBlock(nn.Module):

173 + def __init__(

174 + self,

175 + dim: int,

176 + num_heads: int,

177 + mlp_hidden_dim: int,

178 + act_fn: Callable[[torch.Tensor], torch.Tensor] = F.silu,

179 + norm_layer: Optional[Callable[[int], nn.Module]] = None,

180 + quant_config: Optional[QuantizationConfig] = None,

181 + prefix: str = "",

182 + ) -> None:

183 + super().__init__()

184 + if norm_layer is None:

185 + norm_layer = partial(nn.LayerNorm, eps=1e-6)

186 + self.norm1 = norm_layer(dim)

187 + self.norm2 = norm_layer(dim)

188 + self.attn = Qwen2_5_VisionAttention(

189 + embed_dim=dim,

190 + num_heads=num_heads,

191 + projection_size=dim,

192 + quant_config=quant_config,

193 + prefix=f"{prefix}.attn",

194 + )

195 + self.mlp = Qwen3_VisionMLP(

196 + dim,

197 + mlp_hidden_dim,

198 + act_fn=act_fn,

199 + bias=True,

200 + quant_config=quant_config,

201 + prefix=f"{prefix}.mlp",

202 + )

203 +

204 + def forward(

205 + self,

206 + x: torch.Tensor,

207 + cu_seqlens: torch.Tensor,

208 + rotary_pos_emb: torch.Tensor,

209 + max_seqlen: Optional[int] = None, # Only used for Flash Attention

210 + seqlens: Optional[list[int]] = None, # Only used for xFormers

211 + ) -> torch.Tensor:

212 + x = x + self.attn(

213 + self.norm1(x),

214 + cu_seqlens=cu_seqlens,

215 + rotary_pos_emb=rotary_pos_emb,

216 + max_seqlen=max_seqlen,

217 + seqlens=seqlens,

218 + )

219 +

220 + x = x + self.mlp(self.norm2(x))

221 + return x

222 +

223 +

224 +class Qwen3_VisionPatchMerger(nn.Module):

225 + def __init__(

226 + self,

227 + d_model: int,

228 + context_dim: int,

229 + norm_layer: Optional[Callable[[int], nn.Module]] = None,

230 + spatial_merge_size: int = 2,

231 + use_postshuffle_norm: bool = False,

232 + quant_config: Optional[QuantizationConfig] = None,

233 + prefix: str = "",

234 + ) -> None:

235 + super().__init__()

236 + self.hidden_size = context_dim * (spatial_merge_size**2)

237 +

238 + self.use_postshuffle_norm = use_postshuffle_norm

239 + if self.use_postshuffle_norm:

240 + context_dim = self.hidden_size

241 +

242 + if norm_layer is None:

243 + norm_layer = partial(nn.LayerNorm, eps=1e-6)

244 + self.use_postshuffle_norm = use_postshuffle_norm

245 + self.ln_q = norm_layer(

246 + self.hidden_size if use_postshuffle_norm else context_dim

247 + )

248 + self.mlp = nn.ModuleList(

249 + [

250 + ColumnParallelLinear(

251 + self.hidden_size,

252 + self.hidden_size,

253 + bias=True,

254 + quant_config=quant_config,

255 + prefix=f"{prefix}.mlp.0",

256 + ),

257 + nn.GELU(),

258 + RowParallelLinear(

259 + self.hidden_size,

260 + d_model,

261 + bias=True,

262 + quant_config=quant_config,

263 + prefix=f"{prefix}.mlp.2",

264 + ),

265 + ]

266 + )

267 +

268 + def forward(self, x: torch.Tensor) -> torch.Tensor:

269 + if self.use_postshuffle_norm:

270 + x = self.ln_q(x.view(-1, self.hidden_size))

271 + else:

272 + x = self.ln_q(x).view(-1, self.hidden_size)

273 +

274 + mlp_fc1, mlp_act, mlp_fc2 = self.mlp

275 + x_parallel, _ = mlp_fc1(x)

276 + x_parallel = mlp_act(x_parallel)

277 + out, _ = mlp_fc2(x_parallel)

278 + return out

279 +

280 +

281 +class Qwen3Omni_VisionTransformer(nn.Module):

282 + def __init__(

283 + self,

284 + vision_config,

285 + norm_eps: float = 1e-6,

286 + quant_config: Optional[QuantizationConfig] = None,

287 + prefix: str = "",

288 + ) -> None:

289 + super().__init__()

290 + self.hidden_size = vision_config.hidden_size

291 + self.num_heads = vision_config.num_heads

292 + self.image_size = vision_config.image_size

293 + self.patch_size = vision_config.patch_size

294 + self.spatial_merge_size = vision_config.spatial_merge_size

295 + self.spatial_merge_unit = self.spatial_merge_size**2

296 + self.temporal_patch_size = vision_config.temporal_patch_size

297 + self.num_grid_per_side = self.image_size // self.patch_size

298 + self.apply_vit_abs_pos_embed = vision_config.apply_vit_abs_pos_embed

299 + self.deepstack_visual_indexes = vision_config.deepstack_visual_indexes

300 +

301 + self.patch_embed = Qwen3_VisionPatchEmbed(

302 + patch_size=self.patch_size,

303 + temporal_patch_size=self.temporal_patch_size,

304 + in_channels=vision_config.in_channels,

305 + hidden_size=self.hidden_size,

306 + )

307 +

308 + # vit pos embeding, TODO: spatial_patch_size vs patch_size

309 + if self.apply_vit_abs_pos_embed:

310 + self.pos_embed = nn.Embedding(self.num_grid_per_side**2, self.hidden_size)

311 + else:

312 + self.pos_embed = nn.Parameter(

313 + torch.empty([1, self.num_grid_per_side**2, self.hidden_size])

314 + )

315 +

316 + norm_layer = partial(nn.LayerNorm, eps=norm_eps)

317 + head_dim = self.hidden_size // self.num_heads

318 + self.rotary_pos_emb = Qwen2_5_VisionRotaryEmbedding(head_dim // 2)

319 +

320 + self.blocks = nn.ModuleList(

321 + [

322 + Qwen3_VisionBlock(

323 + dim=self.hidden_size,

324 + num_heads=self.num_heads,

325 + mlp_hidden_dim=vision_config.intermediate_size,

326 + act_fn=_ACTIVATION_REGISTRY[vision_config.hidden_act],

327 + norm_layer=norm_layer,

328 + quant_config=quant_config,

329 + prefix=f"{prefix}.blocks.{layer_idx}",

330 + )

331 + for layer_idx in range(vision_config.depth)

332 + ]

333 + )

334 + self.merger = Qwen3_VisionPatchMerger(

335 + d_model=vision_config.out_hidden_size,

336 + context_dim=self.hidden_size,

337 + norm_layer=norm_layer,

338 + spatial_merge_size=self.spatial_merge_size,

339 + quant_config=quant_config,

340 + prefix=f"{prefix}.merger",

341 + )

342 + if self.deepstack_visual_indexes is not None:

343 + self.merger_list = nn.ModuleList(

344 + [

345 + Qwen3_VisionPatchMerger(

346 + d_model=vision_config.out_hidden_size,

347 + context_dim=self.hidden_size,

348 + spatial_merge_size=self.spatial_merge_size,

349 + use_postshuffle_norm=True,

350 + norm_layer=norm_layer,

351 + quant_config=quant_config,

352 + prefix=f"{prefix}.merger_list.{layer_idx}",

353 + )

354 + for layer_idx in range(len(self.deepstack_visual_indexes))

355 + ]

356 + )

357 +

358 + self.attn_backend = get_vit_attn_backend(

359 + head_size=head_dim, dtype=torch.get_default_dtype()

360 + )

361 + if self.attn_backend != _Backend.FLASH_ATTN and check_upstream_fa_availability(

362 + torch.get_default_dtype()

363 + ):

364 + self.attn_backend = _Backend.FLASH_ATTN

365 +

366 + @property

367 + def dtype(self) -> torch.dtype:

368 + return self.patch_embed.proj.weight.dtype

369 +

370 + @property

371 + def device(self) -> torch.device:

372 + return self.patch_embed.proj.weight.device

373 +

374 + def rot_pos_emb(self, grid_thw):

375 + pos_ids = []

376 + for t, h, w in grid_thw:

377 + hpos_ids = torch.arange(h).unsqueeze(1).expand(-1, w)

378 + hpos_ids = hpos_ids.reshape(

379 + h // self.spatial_merge_size,

380 + self.spatial_merge_size,

381 + w // self.spatial_merge_size,

382 + self.spatial_merge_size,

383 + )

384 + hpos_ids = hpos_ids.permute(0, 2, 1, 3)

385 + hpos_ids = hpos_ids.flatten()

386 +

387 + wpos_ids = torch.arange(w).unsqueeze(0).expand(h, -1)

388 + wpos_ids = wpos_ids.reshape(

389 + h // self.spatial_merge_size,

390 + self.spatial_merge_size,

391 + w // self.spatial_merge_size,

392 + self.spatial_merge_size,

393 + )

394 + wpos_ids = wpos_ids.permute(0, 2, 1, 3)

395 + wpos_ids = wpos_ids.flatten()

396 + pos_ids.append(torch.stack([hpos_ids, wpos_ids], dim=-1).repeat(t, 1))

397 + pos_ids = torch.cat(pos_ids, dim=0)

398 + max_grid_size = grid_thw[:, 1:].max()

399 + rotary_pos_emb_full = self.rotary_pos_emb(max_grid_size)

400 + rotary_pos_emb = rotary_pos_emb_full[pos_ids].flatten(1)

401 + return rotary_pos_emb

402 +

403 + def fast_pos_embed_interpolate(self, grid_thw: list[list[int]]) -> torch.Tensor:

404 + num_grid_per_side = self.num_grid_per_side

405 + m_size = self.spatial_merge_size

406 + hidden_dim = self.pos_embed.embedding_dim

407 +

408 + outputs = []

409 + for t, h, w in grid_thw:

410 + h_idxs = torch.linspace(

411 + 0, num_grid_per_side - 1, h, dtype=torch.float32, device=self.device

412 + )

413 + w_idxs = torch.linspace(

414 + 0, num_grid_per_side - 1, w, dtype=torch.float32, device=self.device

415 + )

416 +

417 + h_floor = h_idxs.to(torch.long)

418 + w_floor = w_idxs.to(torch.long)

419 + h_ceil = torch.clamp(h_floor + 1, max=num_grid_per_side - 1)

420 + w_ceil = torch.clamp(w_floor + 1, max=num_grid_per_side - 1)

421 +

422 + dh = h_idxs - h_floor

423 + dw = w_idxs - w_floor

424 +

425 + # Create meshgrid view for all h, w vars

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

427 + h_floor_grid, w_floor_grid = torch.meshgrid(h_floor, w_floor, indexing="ij")

428 + h_ceil_grid, w_ceil_grid = torch.meshgrid(h_ceil, w_ceil, indexing="ij")

429 + h_floor_grid_idx = h_floor_grid * num_grid_per_side

430 + h_ceil_grid_idx = h_ceil_grid * num_grid_per_side

431 +

432 + # original computation of weights

433 + # w00 = (1 - dh_grid) * (1 - dw_grid)

434 + # w01 = (1 - dh_grid) * dw_grid

435 + # w10 = dh_grid * (1 - dw_grid)

436 + # w11 = dh_grid * dw_grid

437 + # we reuse w11 here to avoid duplicate

438 + # dh_grid * dw_grid computation

439 + w11 = dh_grid * dw_grid

440 + w10 = dh_grid - w11

441 + w01 = dw_grid - w11

442 + w00 = 1 - dh_grid - dw_grid + w11

443 +

444 + idx00 = h_floor_grid_idx + w_floor_grid

445 + idx01 = h_floor_grid_idx + w_ceil_grid

446 + idx10 = h_ceil_grid_idx + w_floor_grid

447 + idx11 = h_ceil_grid_idx + w_ceil_grid

448 +

449 + indices = torch.stack([idx00, idx01, idx10, idx11], dim=0).reshape(4, -1)

450 + weights = torch.stack([w00, w01, w10, w11], dim=0).reshape(4, -1, 1)

451 + weights = weights.to(dtype=self.dtype, device=self.device)

452 +

453 + embeds = self.pos_embed(indices)

454 + weighted_embeds = embeds * weights

455 + p0, p1, p2, p3 = weighted_embeds.unbind(dim=0)

456 + combined = p0 + p1 + p2 + p3

457 +

458 + combined = combined.view(h * w, hidden_dim)

459 + repeated = combined.unsqueeze(0).expand(t, -1, -1).contiguous()

460 + repeated = repeated.view(

461 + t, h // m_size, m_size, w // m_size, m_size, hidden_dim

462 + )

463 + repeated = repeated.permute(0, 1, 3, 2, 4, 5).reshape(-1, hidden_dim)

464 + outputs.append(repeated)

465 +

466 + return torch.cat(outputs, dim=0)

467 +

468 + def compute_attn_mask_seqlen(

469 + self,

470 + cu_seqlens: torch.Tensor,

471 + ) -> tuple[Optional[int], Optional[list[int]]]:

472 + max_seqlen, seqlens = None, None

473 + if self.attn_backend == _Backend.FLASH_ATTN:

474 + max_seqlen = (cu_seqlens[1:] - cu_seqlens[:-1]).max().item()

475 + elif self.attn_backend == _Backend.XFORMERS:

476 + seqlens = (cu_seqlens[1:] - cu_seqlens[:-1]).tolist()

477 + return max_seqlen, seqlens

478 +

479 + def forward(

480 + self,

481 + x: torch.Tensor,

482 + grid_thw: list[list[int]],

483 + ) -> torch.Tensor:

484 + hidden_states = x.to(device=self.device, dtype=self.dtype)

485 + hidden_states = self.patch_embed(hidden_states)

486 +

487 + if self.apply_vit_abs_pos_embed:

488 + pos_embeds = self.fast_pos_embed_interpolate(grid_thw)

489 + hidden_states = hidden_states + pos_embeds

490 + rotary_pos_emb = self.rot_pos_emb(grid_thw)

491 +

492 + cu_seqlens = torch.repeat_interleave(

493 + grid_thw[:, 1] * grid_thw[:, 2], grid_thw[:, 0]

494 + ).cumsum(

495 + dim=0,

496 + dtype=grid_thw.dtype if torch.jit.is_tracing() else torch.int32,

497 + )

498 + cu_seqlens = F.pad(cu_seqlens, (1, 0), value=0)

499 +

500 + hidden_states = hidden_states.unsqueeze(1)

501 + rotary_pos_emb = rotary_pos_emb.to(hidden_states.device)

502 + max_seqlen, seqlens = self.compute_attn_mask_seqlen(cu_seqlens)

503 +

504 + hidden_states_list = []

505 + deepstack_visual_indexes = self.deepstack_visual_indexes

506 +

507 + for layer_num, blk in enumerate(self.blocks):

508 + hidden_states = blk(

509 + hidden_states,

510 + cu_seqlens=cu_seqlens,

511 + rotary_pos_emb=rotary_pos_emb,

512 + max_seqlen=max_seqlen,

513 + seqlens=seqlens,

514 + )

515 + if (

516 + deepstack_visual_indexes is not None

517 + and layer_num in deepstack_visual_indexes

518 + ):

519 + hidden_states_list.append(hidden_states)

520 +

521 + hidden_states = self.merger(hidden_states)

522 +

523 + # processing deepstack

524 + if deepstack_visual_indexes is not None:

525 + processed_hidden_states_list = [hidden_states]

526 + for idx, x in enumerate(hidden_states_list):

527 + x = self.merger_list[idx](x)

528 + processed_hidden_states_list.append(x)

529 + # we cat the original visual features and deepstack features

530 + # along the feature dim

531 + hidden_states = torch.cat(

532 + processed_hidden_states_list, dim=1

533 + ) # [seq_len, hidden_size * (1 + depth_of_deepstack)]

534 +

535 + return hidden_states

536 +

537 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

538 + stacked_params_mapping = [

539 + # (param_name, shard_name, shard_id)

540 + ("attn.qkv.", "attn.q.", "q"),

541 + ("attn.qkv.", "attn.k.", "k"),

542 + ("attn.qkv.", "attn.v.", "v"),

543 + ]

544 + params_dict = dict(self.named_parameters(remove_duplicate=False))

545 + loaded_params: set[str] = set()

546 +

547 + for name, loaded_weight in weights:

548 + for param_name, weight_name, shard_id in stacked_params_mapping:

549 + if weight_name not in name:

550 + continue

551 + name = name.replace(weight_name, param_name)

552 +

553 + param = params_dict[name]

554 + weight_loader = param.weight_loader

555 + weight_loader(param, loaded_weight, shard_id)

556 + break

557 + else:

558 + param = params_dict[name]

559 + weight_loader = getattr(param, "weight_loader", default_weight_loader)

560 + weight_loader(param, loaded_weight)

561 + loaded_params.add(name)

562 + return loaded_params

563 +

564 +

565 +@support_torch_compile(

566 + dynamic_arg_dims={

567 + "input_ids": 0,

568 + "positions": -1,

569 + "intermediate_tensors": 0,

570 + "inputs_embeds": 0,

571 + "deepstack_input_embeds": 0,

572 + }

573 +)

574 +class Qwen3MoeLLMModel(Qwen3MoeModel):

575 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

576 + super().__init__(vllm_config=vllm_config, prefix=prefix)

577 +

578 + self.deepstack_multiscale_layer_start = 1

579 +

580 + def forward(

581 + self,

582 + input_ids: torch.Tensor,

583 + positions: torch.Tensor,

584 + intermediate_tensors: Optional[IntermediateTensors] = None,

585 + inputs_embeds: Optional[torch.Tensor] = None,

586 + deepstack_input_embeds: Optional[IntermediateTensors] = None,

587 + ) -> Union[torch.Tensor, IntermediateTensors]:

588 + if get_pp_group().is_first_rank:

589 + if inputs_embeds is not None:

590 + hidden_states = inputs_embeds

591 + else:

592 + hidden_states = self.get_input_embeddings(input_ids)

593 + residual = None

594 + else:

595 + assert intermediate_tensors is not None

596 + hidden_states = intermediate_tensors["hidden_states"]

597 + residual = intermediate_tensors["residual"]

598 + for layer_idx, layer in enumerate(

599 + self.layers[self.start_layer : self.end_layer]

600 + ):

601 + layer_idx = layer_idx + self.start_layer

602 +

603 + hidden_states, residual = layer(

604 + positions,

605 + hidden_states,

606 + residual,

607 + )

608 +

609 + if deepstack_input_embeds is not None and layer_idx in range(

610 + 0, len(deepstack_input_embeds)

611 + ):

612 + hidden_states = (

613 + hidden_states

614 + + deepstack_input_embeds[f"deepstack_input_embeds_{layer_idx}"]

615 + )

616 +

617 + if not get_pp_group().is_last_rank:

618 + return IntermediateTensors(

619 + {"hidden_states": hidden_states, "residual": residual}

620 + )

621 + hidden_states, _ = self.norm(hidden_states, residual)

622 + return hidden_states

623 +

624 +

625 +class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):

626 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

627 + super(Qwen3MoeForCausalLM, self).__init__()

628 + config = vllm_config.model_config.hf_config

629 + quant_config = vllm_config.quant_config

630 + self.config = config

631 + self.quant_config = quant_config

632 + self.model = Qwen3MoeLLMModel(

633 + vllm_config=vllm_config, prefix=maybe_prefix(prefix, "model")

634 + )

635 + self.lm_head = ParallelLMHead(

636 + config.vocab_size, config.hidden_size, quant_config=quant_config

637 + )

638 + if self.config.tie_word_embeddings:

639 + self.lm_head.weight = self.model.embed_tokens.weight

640 + self.logits_processor = LogitsProcessor(config.vocab_size)

641 + self.make_empty_intermediate_tensors = (

642 + self.model.make_empty_intermediate_tensors

643 + )

644 +

645 +

646 +class Qwen3OmniMoeThinkerProcessingInfo(

647 + Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo

648 +):

649 + def get_hf_config(self):

650 + return self.ctx.get_hf_config(Qwen3OmniMoeConfig).thinker_config

651 +

652 + def get_hf_processor(self, **kwargs: object) -> Qwen3OmniMoeProcessor:

653 + processor = self.ctx.get_hf_processor(

654 + Qwen3OmniMoeProcessor,

655 + use_fast=kwargs.pop("use_fast", True),

656 + **kwargs,

657 + )

658 + if not hasattr(processor, "audio_token"):

659 + processor.audio_token = "<|audio_pad|>"

660 + if not hasattr(processor, "image_token"):

661 + processor.image_token = "<|image_pad|>"

662 + if not hasattr(processor, "video_token"):

663 + processor.video_token = "<|video_pad|>"

664 + return processor

665 +

666 + def get_feature_extractor(self, **kwargs: object):

667 + hf_processor = self.get_hf_processor(**kwargs)

668 + feature_extractor = hf_processor.feature_extractor # type: ignore

669 + assert isinstance(feature_extractor, WhisperFeatureExtractor)

670 + return feature_extractor

671 +

672 + def get_supported_mm_limits(self) -> Mapping[str, Optional[int]]:

673 + return {"audio": None, "image": None, "video": None}

674 +

675 +

676 +Qwen3OmniMoeThinkerDummyInputsBuilder = Qwen2_5OmniThinkerDummyInputsBuilder

677 +

678 +

679 +class Qwen3OmniMoeThinkerMultiModalProcessor(

680 + Qwen2_5OmniThinkerMultiModalProcessor,

681 +):

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 +

692 + def _call_hf_processor(

693 + self,

694 + prompt: str,

695 + mm_data: Mapping[str, object],

696 + mm_kwargs: Mapping[str, object],

697 + tok_kwargs: Mapping[str, object],

698 + ) -> BatchFeature:

699 + mm_data = dict(mm_data)

700 + audios = mm_data.pop("audios", [])

701 +

702 + def pad_to_hop_length(x: np.ndarray, hop_length: int) -> np.ndarray:

703 + length = x.shape[-1]

704 + if length % hop_length != 0:

705 + pad_length = hop_length - (length % hop_length)

706 + x = np.pad(x, (0, pad_length), mode="constant", constant_values=0)

707 + return x

708 +

709 + # NOTE: WhisperFeatureExtractor cannot handle empty list of audios

710 + if audios:

711 + # NOTE: Qwen3-Omni processor accept "audio"

712 + # To make sure the cache works with padding=True, we pre-padded

713 + # the audio to multiple of hop_length.

714 + hop_length = self.info.get_feature_extractor().hop_length

715 + mm_data["audio"] = [

716 + pad_to_hop_length(audio, hop_length)

717 + if isinstance(audio, np.ndarray)

718 + else (pad_to_hop_length(audio[0], hop_length), audio[1])

719 + for audio in audios

720 + ]

721 + mm_kwargs = dict(

722 + **mm_kwargs,

723 + )

724 +

725 + hf_inputs = super()._call_hf_processor(

726 + prompt=prompt,

727 + mm_data=mm_data,

728 + mm_kwargs=mm_kwargs,

729 + tok_kwargs=tok_kwargs,

730 + )

731 +

732 + if (

733 + "audio_feature_lengths" in hf_inputs

734 + and "feature_attention_mask" in hf_inputs

735 + and (audios := mm_data.get("audio", []))

736 + ):

737 + hop_length = self.info.get_feature_extractor().hop_length

738 + audio_num_frames = []

739 + for _, audio in enumerate(audios):

740 + audio_length = len(audio[0]) if isinstance(audio, tuple) else len(audio)

741 + num_frame = (

742 + (audio_length // hop_length)

743 + if audio_length % hop_length == 0

744 + else (audio_length // hop_length - 1)

745 + )

746 + audio_num_frames.append(num_frame)

747 + hf_inputs["feature_attention_mask"] = [

748 + torch.ones(num_frame) for num_frame in audio_num_frames

749 + ]

750 + hf_inputs["audio_feature_lengths"] = torch.tensor(audio_num_frames)

751 + return hf_inputs

752 +

753 + def _maybe_apply_prompt_updates(

754 + self,

755 + mm_items: MultiModalDataItems,

756 + prompt_ids: list[int],

757 + mm_kwargs: MultiModalKwargsItems,

758 + mm_prompt_updates: MultiModalPromptUpdates,

759 + is_update_applied: bool,

760 + ) -> tuple[list[int], str, Mapping[str, list[PlaceholderFeaturesInfo]]]:

761 + """

762 + Qwen3-Omni reimplements this function to handle `use_audio_in_video`.

763 + """

764 + mm_item_counts = mm_items.get_all_counts()

765 + self._validate_mm_kwargs(mm_kwargs, mm_item_counts)

766 +

767 + use_audio_in_video = False

768 + if "video" in mm_kwargs:

769 + for item in mm_kwargs["video"]:

770 + if item and item["use_audio_in_video"].data:

771 + use_audio_in_video = True

772 + else:

773 + use_audio_in_video = False

774 +

775 + if use_audio_in_video and "video" in mm_item_counts:

776 + assert "audio" in mm_item_counts

777 + mm_item_counts["audio"] -= mm_item_counts["video"]

778 +

779 + # Special case with `use_audio_in_video=True`

780 + if use_audio_in_video:

781 + if is_update_applied:

782 + prompt_ids = self._get_raw_input_ids(prompt_ids, use_audio_in_video)

783 + (

784 + prompt_ids,

785 + mm_placeholders,

786 + ) = self._apply_prompt_updates(

787 + prompt_ids,

788 + mm_prompt_updates,

789 + )

790 + self._validate_mm_placeholders(mm_placeholders, mm_item_counts)

791 + # normal case with `use_audio_in_video=False`

792 + elif is_update_applied:

793 + mm_placeholders = self._find_mm_placeholders(

794 + prompt_ids,

795 + mm_prompt_updates,

796 + )

797 + self._validate_mm_placeholders(

798 + mm_placeholders,

799 + mm_item_counts,

800 + )

801 + else:

802 + prompt_ids, mm_placeholders = self._apply_prompt_updates(

803 + prompt_ids,

804 + mm_prompt_updates,

805 + )

806 + self._validate_mm_placeholders(

807 + mm_placeholders,

808 + mm_item_counts,

809 + )

810 +

811 + return prompt_ids, mm_placeholders

812 +

813 + def get_updates_use_audio_in_video(

814 + self,

815 + thinker_config: PretrainedConfig,

816 + audio_len: int,

817 + video_grid_thw: Union[list[int], torch.Tensor],

818 + video_second_per_grid_t: float,

819 + ) -> list[int]:

820 + shift = 0

821 + audio_token_id = thinker_config.audio_token_id

822 + video_token_id = thinker_config.video_token_id

823 + audio_start_token_id = thinker_config.audio_start_token_id

824 + audio_end_token_id = thinker_config.audio_end_token_id

825 + spatial_merge_size = thinker_config.vision_config.spatial_merge_size

826 + position_id_per_seconds = thinker_config.position_id_per_seconds

827 + audio_token_indices = np.arange(next(iter([audio_len])))

828 + curr_video_grid_thw = next(iter([video_grid_thw]))

829 + height = curr_video_grid_thw[1] // spatial_merge_size

830 + width = curr_video_grid_thw[2] // spatial_merge_size

831 + video_token_indices = np.arange(curr_video_grid_thw[0]).reshape(-1, 1, 1)

832 + video_token_indices = np.broadcast_to(

833 + video_token_indices, (video_token_indices.shape[0], height, width)

834 + ).reshape(-1)

835 + video_token_indices = (

836 + (video_token_indices + shift)

837 + * next(iter([video_second_per_grid_t]))

838 + * position_id_per_seconds

839 + )

840 + video_data_index, audio_data_index = 0, 0

841 + updates = [audio_start_token_id]

842 + while video_data_index < len(video_token_indices) and audio_data_index < len(

843 + audio_token_indices

844 + ):

845 + if (

846 + video_token_indices[video_data_index]

847 + <= audio_token_indices[audio_data_index]

848 + ):

849 + updates += [video_token_id]

850 + video_data_index += 1

851 + else:

852 + updates += [audio_token_id]

853 + audio_data_index += 1

854 + if video_data_index < len(video_token_indices):

855 + updates += [video_token_id] * (len(video_token_indices) - video_data_index)

856 + if audio_data_index < len(audio_token_indices):

857 + updates += [audio_token_id] * (len(audio_token_indices) - audio_data_index)

858 + updates += [audio_end_token_id]

859 + return updates

860 +

861 + def _get_prompt_updates(

862 + self,

863 + mm_items: MultiModalDataItems,

864 + hf_processor_mm_kwargs: Mapping[str, Any],

865 + out_mm_kwargs: MultiModalKwargsItems,

866 + ) -> Sequence[PromptUpdate]:

867 + processor = self.info.get_hf_processor(**hf_processor_mm_kwargs)

868 + tokenizer = self.info.get_tokenizer()

869 + image_processor = self.info.get_image_processor(**hf_processor_mm_kwargs)

870 + vocab = tokenizer.get_vocab()

871 +

872 + audio_token = processor.audio_token

873 + image_token = processor.image_token

874 + video_token = processor.video_token

875 + audio_token_id = vocab[audio_token]

876 + image_token_id = vocab[image_token]

877 + video_token_id = vocab[video_token]

878 +

879 + out_mm_data = out_mm_kwargs.get_data()

880 + audio_feature_lengths = out_mm_data.get("audio_feature_lengths")

881 + feature_attention_mask = out_mm_data.get("feature_attention_mask")

882 + if audio_feature_lengths is None and feature_attention_mask is None:

883 + audio_output_lengths = []

884 + elif audio_feature_lengths is not None:

885 + _, audio_output_lens = self._get_feat_extract_output_lengths(

886 + audio_feature_lengths

887 + )

888 + audio_output_lengths = audio_output_lens.tolist()

889 + elif feature_attention_mask is not None:

890 + assert isinstance(feature_attention_mask, torch.Tensor)

891 + _, audio_output_lens = self._get_feat_extract_output_lengths(

892 + feature_attention_mask.sum(-1)

893 + )

894 + audio_output_lengths = audio_output_lens.tolist()

895 +

896 + # number of audios read from video.

897 + audio_in_video_item_idx = 0

898 + audio_item_idx = 0

899 +

900 + def get_replacement_qwen2_audio(item_idx: int):

901 + nonlocal audio_item_idx

902 + item_idx += audio_in_video_item_idx

903 +

904 + audio_item_idx += 1

905 +

906 + num_features = audio_output_lengths[item_idx]

907 + if num_features == 0:

908 + audios = mm_items.get_items("audio", AudioProcessorItems)

909 + audio = audios.get(item_idx)

910 + raise ValueError(

911 + f"The audio {audio} (len={len(audio)}) is too short "

912 + "to be represented inside the model"

913 + )

914 +

915 + return [audio_token_id] * num_features

916 +

917 + def get_replacement_qwen2_vision(item_idx: int, modality: str):

918 + grid_thw = out_mm_data[f"{modality}_grid_thw"][item_idx]

919 + assert isinstance(grid_thw, torch.Tensor)

920 + merge_length = image_processor.merge_size**2

921 +

922 + token_id = image_token_id if modality == "image" else video_token_id

923 + return [token_id] * (int(grid_thw.prod()) // merge_length)

924 +

925 + use_audio_in_video = hf_processor_mm_kwargs.get("use_audio_in_video", False)

926 + thinker_config = self.info.get_hf_config()

927 +

928 + def get_replacement_qwen2_use_audio_in_video(item_idx: int):

929 + nonlocal audio_in_video_item_idx

930 + audio_num_features = audio_output_lengths[audio_item_idx + item_idx]

931 + video_grid_thw = out_mm_data["video_grid_thw"][item_idx]

932 +

933 + audio_in_video_item_idx += 1

934 +

935 + second_per_grid_ts = hf_processor_mm_kwargs.get("second_per_grid_ts", None)

936 + if second_per_grid_ts:

937 + video_second_per_grid_t = second_per_grid_ts[item_idx]

938 + else:

939 + video_second_per_grid_t = 1.0

940 +

941 + return self.get_updates_use_audio_in_video(

942 + thinker_config=thinker_config,

943 + audio_len=audio_num_features,

944 + video_grid_thw=video_grid_thw,

945 + video_second_per_grid_t=video_second_per_grid_t,

946 + )

947 +

948 + video_replacement_fn = (

949 + get_replacement_qwen2_use_audio_in_video

950 + if use_audio_in_video

951 + else partial(get_replacement_qwen2_vision, modality="video")

952 + )

953 +

954 + return [

955 + PromptReplacement(

956 + modality="audio",

957 + target=audio_token,

958 + replacement=get_replacement_qwen2_audio,

959 + ),

960 + PromptReplacement(

961 + modality="image",

962 + target=image_token,

963 + replacement=partial(get_replacement_qwen2_vision, modality="image"),

964 + ),

965 + PromptReplacement(

966 + modality="video",

967 + target=video_token,

968 + replacement=video_replacement_fn,

969 + ),

970 + ]

971 +

972 + def _validate_mm_placeholders(

973 + self,

974 + mm_placeholders: Mapping[str, list[PlaceholderFeaturesInfo]],

975 + mm_item_counts: Mapping[str, int],

976 + ) -> None:

977 + BaseMultiModalProcessor[

978 + Qwen2_5OmniThinkerProcessingInfo

979 + ]._validate_mm_placeholders(self, mm_placeholders, mm_item_counts)

980 +

981 + def _get_raw_input_ids(

982 + self,

983 + token_ids: list[int],

984 + use_audio_in_video: bool = False,

985 + ) -> list[int]:

986 + tokenizer = self.info.get_tokenizer()

987 + vision_bos_token = tokenizer.encode(tokenizer.vision_bos_token)[0]

988 + vision_eos_token = tokenizer.encode(tokenizer.vision_eos_token)[0]

989 + audio_bos_token = tokenizer.encode(tokenizer.audio_bos_token)[0]

990 + audio_eos_token = tokenizer.encode(tokenizer.audio_eos_token)[0]

991 + audio_token = tokenizer.encode("<|audio_pad|>")[0]

992 + image_token = tokenizer.encode("<|image_pad|>")[0]

993 + video_token = tokenizer.encode("<|video_pad|>")[0]

994 +

995 + result = token_ids[:]

996 + if use_audio_in_video:

997 + while True:

998 + start = None

999 + for i in range(len(result) - 1):

1000 + if result[i : i + 2] == [vision_bos_token, audio_bos_token]:

1001 + start = i

1002 + break

1003 + if start is not None:

1004 + end = None

1005 + for i in range(start + 2, len(result) - 1):

1006 + if result[i : i + 2] == [audio_eos_token, vision_eos_token]:

1007 + end = i

1008 + break

1009 + if end is not None:

1010 + result = (

1011 + result[:start]

1012 + + [vision_bos_token, video_token, vision_eos_token]

1013 + + result[end + 2 :]

1014 + )

1015 + else:

1016 + break

1017 +

1018 + for mm_token in [audio_token, image_token, video_token]:

1019 + compressed = []

1020 + for x in result:

1021 + if x != mm_token or (not compressed or compressed[-1] != mm_token):

1022 + compressed.append(x)

1023 + result = compressed

1024 +

1025 + return result

1026 +

1027 +

1028 +class Qwen3OmniMoeConditionalGenerationMixin(Qwen2_5OmniConditionalGenerationMixin):

1029 + def _validate_and_reshape_mm_tensor(

1030 + self, mm_input: object, name: str, dim: int = 0

1031 + ) -> torch.Tensor:

1032 + if not isinstance(mm_input, (torch.Tensor, list)):

1033 + raise ValueError(f"Incorrect type of {name}. Got type: {type(mm_input)}")

1034 + if name == "feature_attention_mask":

1035 + dim = -1

1036 + if isinstance(mm_input, torch.Tensor):

1037 + return torch.concat(list(mm_input), dim=dim)

1038 + else:

1039 + if isinstance(mm_input[0], list):

1040 + return torch.concat(

1041 + [torch.concat(mm_input[i], dim=dim) for i in range(len(mm_input))],

1042 + dim=dim,

1043 + )

1044 + else:

1045 + return torch.concat(mm_input, dim=dim)

1046 +

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 +

1057 + def _process_audio_input(

1058 + self,

1059 + audio_input: Qwen2AudioFeatureInputs,

1060 + audio_hashes: list[str] = None,

1061 + cached_audio_features: torch.Tensor = None,

1062 + ) -> torch.Tensor:

1063 + input_features = audio_input["input_features"]

1064 + audio_feature_lengths = audio_input["audio_feature_lengths"]

1065 +

1066 + if input_features.ndim == 3:

1067 + assert input_features.shape[0] == 1

1068 + input_features = input_features.squeeze(0)

1069 +

1070 + if not isinstance(audio_feature_lengths, torch.Tensor):

1071 + audio_feature_lengths = torch.cat(audio_feature_lengths)

1072 + if audio_feature_lengths.ndim == 2:

1073 + audio_feature_lengths = audio_feature_lengths.reshape(-1)

1074 +

1075 + audio_feat_lengths, audio_output_lengths = (

1076 + self._get_feat_extract_output_lengths(audio_feature_lengths)

1077 + )

1078 +

1079 + audio_outputs = self.audio_tower(

1080 + input_features.to(self.audio_tower.dtype),

1081 + feature_lens=audio_feature_lengths,

1082 + aftercnn_lens=audio_feat_lengths,

1083 + )

1084 + audio_features = audio_outputs.last_hidden_state

1085 + return audio_features.split(audio_output_lengths.tolist())

1086 +

1087 +

1088 +@MULTIMODAL_REGISTRY.register_processor(

1089 + Qwen3OmniMoeThinkerMultiModalProcessor,

1090 + info=Qwen3OmniMoeThinkerProcessingInfo,

1091 + dummy_inputs=Qwen3OmniMoeThinkerDummyInputsBuilder,

1092 +)

1093 +class Qwen3OmniMoeThinkerForConditionalGeneration(

1094 + nn.Module,

1095 + SupportsMultiModal,

1096 + SupportsPP,

1097 + Qwen3OmniMoeConditionalGenerationMixin,

1098 +):

1099 + hf_to_vllm_mapper = WeightsMapper(

1100 + orig_to_new_prefix={

1101 + "thinker.lm_head.": "language_model.lm_head.",

1102 + "thinker.model.": "language_model.model.",

1103 + "thinker.": "",

1104 + }

1105 + )

1106 +

1107 + @classmethod

1108 + def get_placeholder_str(cls, modality: str, i: int) -> Optional[str]:

1109 + if modality.startswith("image"):

1110 + return "<|vision_start|><|image_pad|><|vision_end|>"

1111 + if modality.startswith("video"):

1112 + return "<|vision_start|><|video_pad|><|vision_end|>"

1113 + if modality.startswith("audio"):

1114 + return "<|audio_start|><|audio_pad|><|audio_end|>"

1115 +

1116 + raise ValueError("Only image, video or audio modality is supported")

1117 +

1118 + def __init__(self, *, vllm_config: VllmConfig, prefix: str = ""):

1119 + super().__init__()

1120 + thinker_config: Qwen3OmniMoeThinkerConfig = (

1121 + vllm_config.model_config.hf_config.thinker_config

1122 + )

1123 + quant_config = vllm_config.quant_config

1124 + multimodal_config = vllm_config.model_config.multimodal_config

1125 + self.config = thinker_config

1126 + self.multimodal_config = multimodal_config

1127 +

1128 + # force "use_flash_attention_2=True" to audio tower to align

1129 + # the results.

1130 + if flash_attn is not None:

1131 + audio_config = thinker_config.audio_config

1132 + audio_config._attn_implementation_autoset = True

1133 + audio_config._attn_implementation = "flash_attention_2"

1134 + else:

1135 + logger.warning(

1136 + "flash_attn is not available, the model may not yield the "

1137 + "exactly same result as the transformers implementation "

1138 + "in the audio tower part."

1139 + )

1140 +

1141 + self.audio_tower = Qwen3OmniMoeAudioEncoder(thinker_config.audio_config)

1142 +

1143 + self.visual = Qwen3Omni_VisionTransformer(

1144 + vision_config=thinker_config.vision_config,

1145 + norm_eps=getattr(thinker_config.text_config, "rms_norm_eps", 1e-6),

1146 + quant_config=quant_config,

1147 + prefix=maybe_prefix(prefix, "visual"),

1148 + )

1149 + self.quant_config = quant_config

1150 +

1151 + self.language_model = Qwen3MoeLLMForCausalLM(

1152 + vllm_config=vllm_config.with_hf_config(

1153 + thinker_config.text_config, architectures=["Qwen3MoeForCausalLM"]

1154 + ),

1155 + prefix=maybe_prefix(prefix, "language_model"),

1156 + )

1157 +

1158 + self.make_empty_intermediate_tensors = (

1159 + self.language_model.make_empty_intermediate_tensors

1160 + )

1161 +

1162 + self.use_deepstack = hasattr(

1163 + thinker_config.vision_config, "deepstack_visual_indexes"

1164 + )

1165 + self.deepstack_num_level = (

1166 + len(thinker_config.vision_config.deepstack_visual_indexes)

1167 + if self.use_deepstack

1168 + else 0

1169 + )

1170 + # register buffer for deepstack

1171 + self.deepstack_input_embeds = (

1172 + [

1173 + torch.zeros(

1174 + vllm_config.scheduler_config.max_num_batched_tokens,

1175 + thinker_config.text_config.hidden_size,

1176 + )

1177 + for _ in range(self.deepstack_num_level)

1178 + ]

1179 + if self.use_deepstack

1180 + else None

1181 + )

1182 + self.visual_dim = thinker_config.vision_config.out_hidden_size

1183 + self.multiscale_dim = self.visual_dim * self.deepstack_num_level

1184 +

1185 + def _get_deepstack_input_embeds(self, num_tokens: int) -> IntermediateTensors:

1186 + # get deepstack_input_embeds from buffer, and clear the buffer

1187 + return IntermediateTensors(

1188 + {

1189 + f"deepstack_input_embeds_{idx}": self.deepstack_input_embeds[idx][

1190 + :num_tokens

1191 + ]

1192 + for idx in range(self.deepstack_num_level)

1193 + }

1194 + )

1195 +

1196 + def _set_deepstack_input_embeds(self, deepstack_input_embeds: torch.Tensor) -> None:

1197 + # set deepstack_input_embeds to buffer

1198 + num_tokens = deepstack_input_embeds.size(1)

1199 + if num_tokens > self.deepstack_input_embeds[0].size(0):

1200 + self.deepstack_input_embeds = [

1201 + torch.zeros(

1202 + num_tokens,

1203 + self.config.text_config.hidden_size,

1204 + device=self.deepstack_input_embeds[0].device,

1205 + dtype=self.deepstack_input_embeds[0].dtype,

1206 + )

1207 + for _ in range(self.deepstack_num_level)

1208 + ]

1209 + for idx in range(self.deepstack_num_level):

1210 + self.deepstack_input_embeds[idx][:num_tokens].copy_(

1211 + deepstack_input_embeds[idx]

1212 + )

1213 +

1214 + def _clear_deepstack_input_embeds(self, num_tokens: int) -> None:

1215 + # clear deepstack_input_embeds in buffer

1216 + if num_tokens > 0:

1217 + for idx in range(self.deepstack_num_level):

1218 + self.deepstack_input_embeds[idx][:num_tokens].zero_()

1219 +

1220 + def _parse_and_validate_multimodal_inputs(self, **kwargs: object) -> dict:

1221 + mm_input_by_modality = {}

1222 +

1223 + # Preserve the order of modalities if there are multiple of them

1224 + # from the order of kwargs.

1225 + for input_key in kwargs:

1226 + if (

1227 + input_key in ("pixel_values", "image_embeds")

1228 + and "image" not in mm_input_by_modality

1229 + ):

1230 + mm_input_by_modality["image"] = self._parse_and_validate_image_input(

1231 + **kwargs

1232 + )

1233 + if (

1234 + input_key in ("pixel_values_videos", "video_embeds")

1235 + and "video" not in mm_input_by_modality

1236 + ):

1237 + mm_input_by_modality["video"] = self._parse_and_validate_video_input(

1238 + **kwargs

1239 + )

1240 + if (

1241 + input_key in ("input_audio_features")

1242 + and "audio" not in mm_input_by_modality

1243 + ):

1244 + mm_input_by_modality["audio"] = self._parse_and_validate_audio_input(

1245 + **kwargs

1246 + )

1247 + return mm_input_by_modality

1248 +

1249 + def get_language_model(self) -> torch.nn.Module:

1250 + return self.language_model

1251 +

1252 + def get_multimodal_embeddings(

1253 + self, **kwargs: object

1254 + ) -> Optional[MultiModalEmbeddings]:

1255 + mm_input_by_modality = self._parse_and_validate_multimodal_inputs(**kwargs)

1256 + if not mm_input_by_modality:

1257 + return []

1258 +

1259 + # The result multimodal_embeddings is tuple of tensors, with each

1260 + # tensor correspoending to a multimodal data item (image or video).

1261 + multimodal_embeddings: tuple[torch.Tensor, ...] = ()

1262 +

1263 + # NOTE: It is important to iterate over the keys in this dictionary

1264 + # to preserve the order of the modalities.

1265 + for modality in mm_input_by_modality:

1266 + multimodal_input = mm_input_by_modality[modality]

1267 + if modality == "image":

1268 + vision_embeddings = self._process_image_input(multimodal_input)

1269 + multimodal_embeddings += vision_embeddings

1270 + if modality == "video":

1271 + video_embeddings = self._process_video_input(multimodal_input)

1272 + multimodal_embeddings += video_embeddings

1273 + if modality == "audio":

1274 + audio_embeddings = self._process_audio_input(multimodal_input)

1275 + multimodal_embeddings += audio_embeddings

1276 + return multimodal_embeddings

1277 +

1278 + def get_input_embeddings(

1279 + self,

1280 + input_ids: torch.Tensor,

1281 + multimodal_embeddings: Optional[MultiModalEmbeddings] = None,

1282 + *,

1283 + is_multimodal: Optional[torch.Tensor] = None,

1284 + handle_oov_mm_token: bool = False,

1285 + ) -> torch.Tensor:

1286 + inputs_embeds = self._get_text_embeddings(

1287 + input_ids,

1288 + self.language_model.get_input_embeddings,

1289 + is_multimodal=is_multimodal,

1290 + handle_oov_mm_token=handle_oov_mm_token,

1291 + )

1292 +

1293 + if multimodal_embeddings is None or len(multimodal_embeddings) == 0:

1294 + return inputs_embeds

1295 +

1296 + deepstack_input_embeds = None

1297 + # TODO (ywang96): support overlapping modalitiy embeddings so that

1298 + # `use_audio_in_video` will work on V1.

1299 + # split the feat dim to obtain multi-scale visual feature

1300 + has_vision_embeddings = [

1301 + embeddings.shape[-1] != self.config.text_config.hidden_size

1302 + for embeddings in multimodal_embeddings

1303 + ]

1304 + if self.visual.deepstack_visual_indexes is not None and any(

1305 + has_vision_embeddings

1306 + ):

1307 + multiscale_len = len(self.visual.deepstack_visual_indexes)

1308 + multimodal_embeddings_multiscale = []

1309 + is_vision = torch.zeros_like(is_multimodal)

1310 + mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]

1311 + mm_position_idx = 0

1312 + for index, embeddings in enumerate(multimodal_embeddings):

1313 + num_tokens = embeddings.shape[0]

1314 + current_positions = mm_positions[

1315 + mm_position_idx : mm_position_idx + num_tokens

1316 + ]

1317 +

1318 + # Vision embeddings

1319 + if embeddings.shape[-1] != self.config.text_config.hidden_size:

1320 + visual_dim = embeddings.shape[-1] // (multiscale_len + 1)

1321 + multi_dim = visual_dim * multiscale_len

1322 + embeddings_main, embeddings_multiscale = torch.split(

1323 + embeddings, [visual_dim, multi_dim], dim=-1

1324 + )

1325 + multimodal_embeddings[index] = embeddings_main

1326 + multimodal_embeddings_multiscale.append(embeddings_multiscale)

1327 + is_vision[current_positions] = True

1328 +

1329 + # Audio embeddings

1330 + else:

1331 + is_vision[current_positions] = False

1332 +

1333 + mm_position_idx += num_tokens

1334 +

1335 + deepstack_input_embeds = inputs_embeds.new_zeros(

1336 + inputs_embeds.size(0), multiscale_len * inputs_embeds.size(1)

1337 + )

1338 + deepstack_input_embeds = _merge_multimodal_embeddings(

1339 + inputs_embeds=deepstack_input_embeds,

1340 + multimodal_embeddings=multimodal_embeddings_multiscale,

1341 + is_multimodal=is_vision,

1342 + )

1343 + deepstack_input_embeds = (

1344 + deepstack_input_embeds.view(

1345 + inputs_embeds.shape[0], multiscale_len, visual_dim

1346 + )

1347 + .permute(1, 0, 2)

1348 + .contiguous()

1349 + )

1350 + self._set_deepstack_input_embeds(deepstack_input_embeds)

1351 +

1352 + inputs_embeds = _merge_multimodal_embeddings(

1353 + inputs_embeds=inputs_embeds,

1354 + multimodal_embeddings=multimodal_embeddings,

1355 + is_multimodal=is_multimodal,

1356 + )

1357 +

1358 + return inputs_embeds

1359 +

1360 + def forward(

1361 + self,

1362 + input_ids: torch.Tensor,

1363 + positions: torch.Tensor,

1364 + intermediate_tensors: Optional[IntermediateTensors] = None,

1365 + inputs_embeds: Optional[torch.Tensor] = None,

1366 + **kwargs: object,

1367 + ) -> Union[torch.Tensor, IntermediateTensors]:

1368 + if intermediate_tensors is not None:

1369 + inputs_embeds = None

1370 +

1371 + if (

1372 + self.use_deepstack

1373 + and inputs_embeds is not None

1374 + and get_pp_group().is_first_rank

1375 + ):

1376 + deepstack_input_embeds = self._get_deepstack_input_embeds(

1377 + inputs_embeds.size(0)

1378 + )

1379 + else:

1380 + deepstack_input_embeds = None

1381 +

1382 + hidden_states = self.language_model.model(

1383 + input_ids,

1384 + positions,

1385 + intermediate_tensors,

1386 + inputs_embeds=inputs_embeds,

1387 + # args for deepstack

1388 + deepstack_input_embeds=deepstack_input_embeds,

1389 + )

1390 +

1391 + if inputs_embeds is not None and get_pp_group().is_first_rank:

1392 + self._clear_deepstack_input_embeds(inputs_embeds.size(0))

1393 +

1394 + return hidden_states

1395 +

1396 + def compute_logits(

1397 + self,

1398 + hidden_states: torch.Tensor,

1399 + ) -> Optional[torch.Tensor]:

1400 + return self.language_model.compute_logits(hidden_states)

1401 +

1402 + def load_weights(self, weights: Iterable[tuple[str, torch.Tensor]]) -> set[str]:

1403 + loader = AutoWeightsLoader(

1404 + self,

1405 + skip_prefixes=["talker.", "code2wav."],

1406 + )

1407 + loaded_weights = loader.load_weights(weights, mapper=self.hf_to_vllm_mapper)

1408 +

1409 + return loaded_weights

vllm/model_executor/models/registry.py

@@ -355,6 +355,10 @@ _MULTIMODAL_MODELS = {

355 "qwen2_5_omni_thinker",

356 "Qwen2_5OmniThinkerForConditionalGeneration",

357 ),

358 + "Qwen3OmniMoeForConditionalGeneration": (

359 + "qwen3_omni_moe_thinker",

360 + "Qwen3OmniMoeThinkerForConditionalGeneration",

361 + ),

362 "Qwen3VLForConditionalGeneration": ("qwen3_vl", "Qwen3VLForConditionalGeneration"), # noqa: E501

363 "Qwen3VLMoeForConditionalGeneration": (

364 "qwen3_vl_moe",