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 9, 14:45 UTC
[Model] Use mm_position to compute mrope positions for Qwen3-Omni (#33010)
Signed-off-by: Itay Etelis <itay.etelis@ibm.com> Co-authored-by: Itay Etelis <itay.etelis@ibm.com>
Itay Etelis committed
Jan 26, 2026 at 15:48 UTC
6ca2c91b9663466b61ab14d804ef3f6fe7c80827
2 files changed
+293
-298
examples/offline_inference/qwen3_omni/only_thinker.py
+56
-3
@@ -2,7 +2,7 @@
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
"""
4
This example shows how to use vLLM for running offline inference
5
-with the correct prompt format on Qwen2.5-Omni (thinker only).
5
+with the correct prompt format on Qwen3-Omni (thinker only).
6
"""
7
8
from typing import NamedTuple
@@ -112,23 +112,51 @@ def get_multi_audios_query() -> QueryResult:
112
)
113
114
115
+def get_multi_images_query() -> QueryResult:
116
+ question = "What are the differences between these two images?"
117
+ prompt = (
118
+ f"<|im_start|>system\n{default_system}<|im_end|>\n"
119
+ "<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>"
120
+ "<|vision_start|><|image_pad|><|vision_end|>"
121
+ f"{question}<|im_end|>\n"
122
+ f"<|im_start|>assistant\n"
123
+ )
124
+ return QueryResult(
125
+ inputs={
126
+ "prompt": prompt,
127
+ "multi_modal_data": {
128
+ "image": [
129
+ convert_image_mode(ImageAsset("cherry_blossom").pil_image, "RGB"),
130
+ convert_image_mode(ImageAsset("stop_sign").pil_image, "RGB"),
131
+ ],
132
+ },
133
+ },
134
+ limit_mm_per_prompt={
135
+ "image": 2,
136
+ },
137
+ )
138
+
139
+
140
query_map = {
141
"mixed_modalities": get_mixed_modalities_query,
142
"use_audio_in_video": get_use_audio_in_video_query,
143
"multi_audios": get_multi_audios_query,
144
+ "multi_images": get_multi_images_query,
145
}
146
147
148
def main(args):
123
- model_name = "Qwen/Qwen3-Omni-30B-A3B-Instruct"
149
+ model_name = args.model
150
query_result = query_map[args.query_type]()
151
152
llm = LLM(
153
model=model_name,
128
- max_model_len=12800,
154
+ max_model_len=args.max_model_len,
155
max_num_seqs=5,
156
limit_mm_per_prompt=query_result.limit_mm_per_prompt,
157
seed=args.seed,
158
+ tensor_parallel_size=args.tensor_parallel_size,
159
+ gpu_memory_utilization=args.gpu_memory_utilization,
160
)
161
162
# We set temperature to 0.2 so that outputs can be different
@@ -161,6 +189,31 @@ def parse_args():
189
default=0,
190
help="Set the seed when initializing `vllm.LLM`.",
191
)
192
+ parser.add_argument(
193
+ "--model",
194
+ type=str,
195
+ default="Qwen/Qwen3-Omni-30B-A3B-Instruct",
196
+ help="Model name or path.",
197
+ )
198
+ parser.add_argument(
199
+ "--tensor-parallel-size",
200
+ "-tp",
201
+ type=int,
202
+ default=1,
203
+ help="Tensor parallel size for distributed inference.",
204
+ )
205
+ parser.add_argument(
206
+ "--gpu-memory-utilization",
207
+ type=float,
208
+ default=0.9,
209
+ help="GPU memory utilization (0.0 to 1.0).",
210
+ )
211
+ parser.add_argument(
212
+ "--max-model-len",
213
+ type=int,
214
+ default=12800,
215
+ help="Maximum model context length.",
216
+ )
217
218
return parser.parse_args()
219
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+237
-295
@@ -22,7 +22,7 @@
22
# limitations under the License.
23
"""Inference-only Qwen3-Omni-Moe model (thinker part)."""
24
25
-from collections.abc import Callable, Iterable, Mapping, Sequence
25
+from collections.abc import Callable, Iterable, Iterator, Mapping, Sequence
26
from functools import partial
27
from typing import Any
28
@@ -104,10 +104,7 @@ from .utils import (
104
_merge_multimodal_embeddings,
105
maybe_prefix,
106
)
107
-from .vision import (
108
- get_llm_pos_ids_for_vision,
109
- get_vit_attn_backend,
110
-)
107
+from .vision import get_vit_attn_backend
108
109
logger = init_logger(__name__)
110
@@ -1867,323 +1864,268 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1864
1865
return loaded_weights
1866
1870
- def get_mrope_input_positions(
1871
- self,
1872
- input_tokens: list[int],
1873
- mm_features: list[MultiModalFeatureSpec],
1874
- ) -> tuple[torch.Tensor, int]:
1875
- kwargs = MultiModalFeatureSpec.gather_kwargs(
1876
- mm_features,
1877
- {
1878
- "image_grid_thw",
1879
- "video_grid_thw",
1880
- "second_per_grid_ts",
1881
- "audio_feature_lengths",
1882
- "use_audio_in_video",
1883
- },
1884
- )
1885
- image_grid_thw = kwargs.get("image_grid_thw", [])
1886
- video_grid_thw = kwargs.get("video_grid_thw", [])
1887
- second_per_grid_ts = kwargs.get("second_per_grid_ts", [])
1888
- audio_feature_lengths = kwargs.get("audio_feature_lengths", [])
1889
- use_audio_in_video = any(kwargs.get("use_audio_in_video", []))
1890
-
1891
- image_grid_thw = (torch.stack if image_grid_thw else torch.tensor)(
1892
- image_grid_thw
1893
- )
1894
- video_grid_thw = (torch.stack if video_grid_thw else torch.tensor)(
1895
- video_grid_thw
1896
- )
1897
-
1898
- input_ids = torch.tensor(input_tokens)
1899
- if input_ids is None or input_ids.ndim != 1:
1900
- raise ValueError("_omni3_get_input_positions_tensor expects 1D input_ids")
1867
+ def _compute_audio_token_count(self, audio_feature_length: int) -> int:
1868
+ """Compute audio tokens from feature length using Qwen3-Omni formula."""
1869
+ return _get_feat_extract_output_lengths(
1870
+ torch.tensor([audio_feature_length])
1871
+ ).item()
1872
1902
- seq_len = input_ids.shape[0]
1873
+ def _get_audio_for_video_mapping(
1874
+ self, mm_features: list[MultiModalFeatureSpec]
1875
+ ) -> tuple[dict[int, int], set[int]]:
1876
+ """
1877
+ Map video offset -> paired audio_feature_length for use_audio_in_video.
1878
1904
- if isinstance(audio_feature_lengths, list):
1905
- audio_feature_lengths = torch.tensor(
1906
- audio_feature_lengths, dtype=torch.long
1907
- )
1879
+ When use_audio_in_video=True, audio is interleaved within video.
1880
+ The pairing is based on feature order in mm_features.
1881
1909
- if not len(second_per_grid_ts) and len(video_grid_thw):
1910
- second_per_grid_ts = 2.0
1911
- second_per_grids = (
1912
- torch.ones(len(video_grid_thw), dtype=torch.float32)
1913
- * second_per_grid_ts
1914
- )
1915
- else:
1916
- second_per_grids = torch.tensor(second_per_grid_ts, dtype=torch.float32)
1882
+ Returns:
1883
+ Tuple of (video_offset -> audio_feature_length mapping,
1884
+ set of paired audio offsets to skip)
1885
+ """
1886
+ videos_with_audio = [
1887
+ f
1888
+ for f in mm_features
1889
+ if f.modality == "video"
1890
+ and f.data.get("use_audio_in_video")
1891
+ and f.data["use_audio_in_video"].data.item()
1892
+ ]
1893
+ audios = [f for f in mm_features if f.modality == "audio"]
1894
+
1895
+ mapping: dict[int, int] = {}
1896
+ paired_audio_offsets: set[int] = set()
1897
+ for i, video_f in enumerate(videos_with_audio):
1898
+ if i < len(audios):
1899
+ audio_len = audios[i].data["audio_feature_lengths"].data.item()
1900
+ mapping[video_f.mm_position.offset] = audio_len
1901
+ paired_audio_offsets.add(audios[i].mm_position.offset)
1902
+ return mapping, paired_audio_offsets
1903
+
1904
+ def iter_mm_features(
1905
+ self, mm_features: list[MultiModalFeatureSpec]
1906
+ ) -> Iterator[tuple[int, str, dict[str, Any]]]:
1907
+ """
1908
+ Iterate over multimodal features sorted by position offset.
1909
1910
+ Yields: (offset, modality, feature_data) where feature_data contains:
1911
+ - image: {"grid_t", "grid_h", "grid_w", "t_factor"}
1912
+ - video: {"grid_t", "grid_h", "grid_w", "t_factor",
1913
+ "use_audio_in_video", "audio_feature_length"}
1914
+ - audio: {"audio_feature_length"}
1915
+ """
1916
config = self.config
1917
spatial_merge_size = config.vision_config.spatial_merge_size
1920
- image_token_id = config.image_token_id
1921
- video_token_id = config.video_token_id
1922
- audio_token_id = config.audio_token_id
1923
- vision_start_token_id = config.vision_start_token_id
1924
- audio_start_token_id = config.audio_start_token_id
1918
position_id_per_seconds = config.position_id_per_seconds
1919
1927
- vision_start_indices = torch.argwhere(
1928
- input_ids == vision_start_token_id
1929
- ).squeeze(1)
1930
- if vision_start_indices.numel() > 0:
1931
- vision_tokens = input_ids[vision_start_indices + 1]
1932
- else:
1933
- vision_tokens = input_ids.new_empty((0,), dtype=input_ids.dtype)
1934
- audio_nums = torch.sum(input_ids == audio_start_token_id)
1935
- image_nums = (vision_tokens == image_token_id).sum()
1936
- video_nums = (
1937
- (vision_tokens == audio_start_token_id).sum()
1938
- if use_audio_in_video
1939
- else (vision_tokens == video_token_id).sum()
1920
+ sorted_features = sorted(mm_features, key=lambda f: f.mm_position.offset)
1921
+ audio_for_video, paired_audio_offsets = self._get_audio_for_video_mapping(
1922
+ sorted_features
1923
)
1924
1942
- llm_pos_ids_list: list[torch.Tensor] = []
1943
- st = 0
1944
- image_idx = 0
1945
- video_idx = 0
1946
- audio_idx = 0
1947
- remain_images, remain_videos, remain_audios = image_nums, video_nums, audio_nums # noqa: E501
1948
- multimodal_nums = (
1949
- image_nums + audio_nums
1950
- if use_audio_in_video
1951
- else image_nums + video_nums + audio_nums
1952
- ) # noqa: E501
1925
+ for mm_feature in sorted_features:
1926
+ offset = mm_feature.mm_position.offset
1927
+ modality = mm_feature.modality
1928
1954
- for _ in range(multimodal_nums):
1955
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1956
- if (image_token_id in input_tokens or video_token_id in input_tokens) and (
1957
- remain_videos > 0 or remain_images > 0
1958
- ):
1959
- ed_vision_start = input_tokens.index(vision_start_token_id, st)
1960
- else:
1961
- ed_vision_start = len(input_tokens) + 1
1962
- if audio_token_id in input_tokens and remain_audios > 0:
1963
- ed_audio_start = input_tokens.index(audio_start_token_id, st)
1964
- else:
1965
- ed_audio_start = len(input_tokens) + 1
1966
- min_ed = min(ed_vision_start, ed_audio_start)
1967
-
1968
- if min_ed == ed_audio_start:
1969
- text_len = min_ed - st
1970
- if text_len != 0:
1971
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1972
- llm_pos_ids_list.append(
1973
- torch.arange(text_len, dtype=torch.long)
1974
- .view(1, -1)
1975
- .expand(3, -1)
1976
- + st_idx
1977
- )
1978
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1979
- bos_len = 1
1980
- llm_pos_ids_list.append(
1981
- torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1982
- + st_idx
1983
- )
1984
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1985
- audio_len = _get_feat_extract_output_lengths(
1986
- audio_feature_lengths[audio_idx]
1929
+ if modality == "image":
1930
+ t, h, w = mm_feature.data["image_grid_thw"].data.tolist()
1931
+ yield (
1932
+ offset,
1933
+ "image",
1934
+ {
1935
+ "grid_t": t,
1936
+ "grid_h": h // spatial_merge_size,
1937
+ "grid_w": w // spatial_merge_size,
1938
+ "t_factor": position_id_per_seconds,
1939
+ },
1940
)
1988
- llm_pos_ids = (
1989
- torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1)
1990
- + st_idx
1941
+ elif modality == "video":
1942
+ t, h, w = mm_feature.data["video_grid_thw"].data.tolist()
1943
+ second_per_grid_ts = 2.0
1944
+ if mm_feature.data.get("second_per_grid_ts"):
1945
+ second_per_grid_ts = mm_feature.data[
1946
+ "second_per_grid_ts"
1947
+ ].data.item()
1948
+ use_audio_in_video = bool(
1949
+ mm_feature.data.get("use_audio_in_video")
1950
+ and mm_feature.data["use_audio_in_video"].data.item()
1951
)
1992
- llm_pos_ids_list.append(llm_pos_ids)
1993
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
1994
- eos_len = 1
1995
- llm_pos_ids_list.append(
1996
- torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
1997
- + st_idx
1952
+
1953
+ yield (
1954
+ offset,
1955
+ "video",
1956
+ {
1957
+ "grid_t": t,
1958
+ "grid_h": h // spatial_merge_size,
1959
+ "grid_w": w // spatial_merge_size,
1960
+ "t_factor": second_per_grid_ts * position_id_per_seconds,
1961
+ "use_audio_in_video": use_audio_in_video,
1962
+ "audio_feature_length": audio_for_video.get(offset),
1963
+ },
1964
)
1999
- st += text_len + bos_len + audio_len + eos_len
1965
+ elif modality == "audio":
1966
+ if offset not in paired_audio_offsets:
1967
+ audio_len = mm_feature.data["audio_feature_lengths"].data.item()
1968
+ yield offset, "audio", {"audio_feature_length": audio_len}
1969
+
1970
+ def _compute_interleaved_positions(
1971
+ self, start_idx: int, data: dict[str, Any]
1972
+ ) -> tuple[np.ndarray, int]:
1973
+ """
1974
+ Compute positions for interleaved video+audio using Qwen3 token-by-token
1975
+ interleaving logic.
1976
+
1977
+ Returns: (position_ids [3, N], total_token_count)
1978
+ """
1979
+ grid_t = data["grid_t"]
1980
+ grid_h = data["grid_h"]
1981
+ grid_w = data["grid_w"]
1982
+ t_factor = data["t_factor"]
1983
+ audio_feature_length = data["audio_feature_length"]
1984
+
1985
+ audio_len = self._compute_audio_token_count(audio_feature_length)
1986
+
1987
+ h_index = np.tile(
1988
+ np.arange(grid_h).reshape(1, -1, 1), (grid_t, 1, grid_w)
1989
+ ).flatten()
1990
+ w_index = np.tile(
1991
+ np.arange(grid_w).reshape(1, 1, -1), (grid_t, grid_h, 1)
1992
+ ).flatten()
1993
+ t_index_raw = np.arange(grid_t)
1994
+ t_index_scaled = (t_index_raw * t_factor).astype(np.int64)
1995
+ t_index = np.repeat(t_index_scaled, grid_h * grid_w)
1996
+
1997
+ video_pos = np.stack([t_index, h_index, w_index]) + start_idx
1998
+ audio_pos = np.broadcast_to(np.arange(audio_len), (3, audio_len)) + start_idx
1999
+
2000
+ video_t_values = video_pos[0]
2001
+ audio_t_values = audio_pos[0]
2002
+
2003
+ pos_ids_list: list[np.ndarray] = []
2004
+ video_idx, audio_idx = 0, 0
2005
+ num_video = grid_t * grid_h * grid_w
2006
+
2007
+ while video_idx < num_video and audio_idx < audio_len:
2008
+ if video_t_values[video_idx] <= audio_t_values[audio_idx]:
2009
+ pos_ids_list.append(video_pos[:, video_idx : video_idx + 1])
2010
+ video_idx += 1
2011
+ else:
2012
+ pos_ids_list.append(audio_pos[:, audio_idx : audio_idx + 1])
2013
audio_idx += 1
2001
- remain_audios -= 1
2002
- elif (
2003
- min_ed == ed_vision_start
2004
- and input_ids[ed_vision_start + 1] == image_token_id
2005
- ):
2006
- text_len = min_ed - st
2007
- if text_len != 0:
2008
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2009
- llm_pos_ids_list.append(
2010
- torch.arange(text_len, dtype=torch.long)
2011
- .view(1, -1)
2012
- .expand(3, -1)
2013
- + st_idx
2014
- )
2015
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2016
- bos_len = 1
2017
- llm_pos_ids_list.append(
2018
- torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
2019
- + st_idx
2020
- )
2021
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2022
- grid_t = image_grid_thw[image_idx][0]
2023
- grid_hs = image_grid_thw[:, 1]
2024
- grid_ws = image_grid_thw[:, 2]
2025
- t_index = torch.arange(grid_t) * position_id_per_seconds
2026
- llm_pos_ids = get_llm_pos_ids_for_vision(
2027
- st_idx, image_idx, spatial_merge_size, t_index, grid_hs, grid_ws
2028
- )
2029
- image_len = image_grid_thw[image_idx].prod() // (spatial_merge_size**2)
2030
- llm_pos_ids_list.append(llm_pos_ids)
2031
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2032
- eos_len = 1
2033
- llm_pos_ids_list.append(
2034
- torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
2035
- + st_idx
2036
- )
2037
- st += text_len + bos_len + image_len + eos_len
2038
- image_idx += 1
2039
- remain_images -= 1
2040
- elif (
2041
- min_ed == ed_vision_start
2042
- and input_ids[ed_vision_start + 1] == video_token_id
2043
- and not use_audio_in_video
2044
- ):
2045
- text_len = min_ed - st
2046
- if text_len != 0:
2047
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2048
- llm_pos_ids_list.append(
2049
- torch.arange(text_len, dtype=torch.long)
2050
- .view(1, -1)
2051
- .expand(3, -1)
2052
- + st_idx
2053
- )
2054
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2055
- bos_len = 1
2056
- llm_pos_ids_list.append(
2057
- torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
2058
- + st_idx
2059
- )
2060
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2061
- grid_t = video_grid_thw[video_idx][0]
2062
- grid_hs = video_grid_thw[:, 1]
2063
- grid_ws = video_grid_thw[:, 2]
2064
- t_index = (
2065
- torch.arange(grid_t)
2066
- * float(second_per_grids[video_idx].item())
2067
- * position_id_per_seconds
2068
- )
2069
- llm_pos_ids = get_llm_pos_ids_for_vision(
2070
- st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
2071
- )
2072
- video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
2073
- llm_pos_ids_list.append(llm_pos_ids)
2074
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2075
- eos_len = 1
2014
+
2015
+ if video_idx < num_video:
2016
+ pos_ids_list.append(video_pos[:, video_idx:])
2017
+ if audio_idx < audio_len:
2018
+ pos_ids_list.append(audio_pos[:, audio_idx:])
2019
+
2020
+ total_tokens = num_video + audio_len
2021
+ return np.concatenate(pos_ids_list, axis=1), total_tokens
2022
+
2023
+ def get_mrope_input_positions(
2024
+ self,
2025
+ input_tokens: list[int],
2026
+ mm_features: list[MultiModalFeatureSpec],
2027
+ ) -> tuple[torch.Tensor, int]:
2028
+ """Compute M-RoPE input positions using mm_features directly."""
2029
+ seq_len = len(input_tokens)
2030
+
2031
+ llm_pos_ids_list: list[np.ndarray] = []
2032
+ st = 0
2033
+
2034
+ for offset, modality, data in self.iter_mm_features(mm_features):
2035
+ text_len = offset - st
2036
+ st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
2037
+
2038
+ if text_len > 0:
2039
llm_pos_ids_list.append(
2077
- torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
2078
- + st_idx
2040
+ np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx
2041
)
2080
- st += text_len + bos_len + video_len + eos_len
2081
- video_idx += 1
2082
- remain_videos -= 1
2083
- elif (
2084
- min_ed == ed_vision_start
2085
- and ed_vision_start + 1 == ed_audio_start
2086
- and use_audio_in_video
2087
- ):
2088
- text_len = min_ed - st
2089
- if text_len != 0:
2090
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2091
- llm_pos_ids_list.append(
2092
- torch.arange(text_len, dtype=torch.long)
2093
- .view(1, -1)
2094
- .expand(3, -1)
2095
- + st_idx
2096
- )
2097
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2098
- bos_len = 1
2099
- bos_block = (
2100
- torch.arange(bos_len, dtype=torch.long).view(1, -1).expand(3, -1)
2101
- + st_idx
2102
- )
2103
- llm_pos_ids_list.append(bos_block)
2104
- llm_pos_ids_list.append(bos_block)
2105
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2106
- audio_len = _get_feat_extract_output_lengths(
2107
- audio_feature_lengths[audio_idx]
2108
- )
2109
- audio_llm_pos_ids = (
2110
- torch.arange(audio_len, dtype=torch.long).view(1, -1).expand(3, -1)
2111
- + st_idx
2112
- )
2113
- grid_t = video_grid_thw[video_idx][0]
2114
- grid_hs = video_grid_thw[:, 1]
2115
- grid_ws = video_grid_thw[:, 2]
2116
- t_index = (
2117
- torch.arange(grid_t)
2118
- * float(second_per_grids[video_idx].item())
2119
- * position_id_per_seconds
2042
+ st_idx += text_len
2043
+
2044
+ bos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
2045
+ llm_pos_ids_list.append(bos_pos)
2046
+ st_idx += 1
2047
+
2048
+ if modality == "audio":
2049
+ audio_tokens = self._compute_audio_token_count(
2050
+ data["audio_feature_length"]
2051
)
2121
- video_llm_pos_ids = get_llm_pos_ids_for_vision(
2122
- st_idx, video_idx, spatial_merge_size, t_index, grid_hs, grid_ws
2052
+ audio_pos = (
2053
+ np.broadcast_to(np.arange(audio_tokens), (3, audio_tokens)) + st_idx
2054
)
2124
- video_data_index, audio_data_index = 0, 0
2125
- while (
2126
- video_data_index < video_llm_pos_ids.shape[-1]
2127
- and audio_data_index < audio_llm_pos_ids.shape[-1]
2128
- ):
2129
- if (
2130
- video_llm_pos_ids[0][video_data_index]
2131
- <= audio_llm_pos_ids[0][audio_data_index]
2132
- ):
2133
- llm_pos_ids_list.append(
2134
- video_llm_pos_ids[
2135
- :, video_data_index : video_data_index + 1
2136
- ]
2137
- )
2138
- video_data_index += 1
2139
- else:
2140
- llm_pos_ids_list.append(
2141
- audio_llm_pos_ids[
2142
- :, audio_data_index : audio_data_index + 1
2143
- ]
2144
- )
2145
- audio_data_index += 1
2146
- if video_data_index < video_llm_pos_ids.shape[-1]:
2147
- llm_pos_ids_list.append(
2148
- video_llm_pos_ids[
2149
- :, video_data_index : video_llm_pos_ids.shape[-1]
2150
- ]
2151
- )
2152
- if audio_data_index < audio_llm_pos_ids.shape[-1]:
2153
- llm_pos_ids_list.append(
2154
- audio_llm_pos_ids[
2155
- :, audio_data_index : audio_llm_pos_ids.shape[-1]
2156
- ]
2055
+ llm_pos_ids_list.append(audio_pos)
2056
+ st_idx = int(audio_pos.max()) + 1
2057
+
2058
+ eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
2059
+ llm_pos_ids_list.append(eos_pos)
2060
+ st = offset + 1 + audio_tokens + 1
2061
+
2062
+ elif modality == "image":
2063
+ grid_t = data["grid_t"]
2064
+ grid_h = data["grid_h"]
2065
+ grid_w = data["grid_w"]
2066
+ t_factor = data["t_factor"]
2067
+
2068
+ grid_indices = np.indices((grid_t, grid_h, grid_w))
2069
+ if t_factor != 1.0:
2070
+ grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
2071
+ llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
2072
+
2073
+ image_len = grid_t * grid_h * grid_w
2074
+ st_idx = int(llm_pos_ids_list[-1].max()) + 1
2075
+
2076
+ eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
2077
+ llm_pos_ids_list.append(eos_pos)
2078
+ st = offset + 1 + image_len + 1
2079
+
2080
+ elif modality == "video":
2081
+ grid_t = data["grid_t"]
2082
+ grid_h = data["grid_h"]
2083
+ grid_w = data["grid_w"]
2084
+ t_factor = data["t_factor"]
2085
+
2086
+ if not data["use_audio_in_video"]:
2087
+ grid_indices = np.indices((grid_t, grid_h, grid_w))
2088
+ if t_factor != 1.0:
2089
+ grid_indices[0] = (grid_indices[0] * t_factor).astype(np.int64)
2090
+ llm_pos_ids_list.append(grid_indices.reshape(3, -1) + st_idx)
2091
+
2092
+ video_len = grid_t * grid_h * grid_w
2093
+ st_idx = int(llm_pos_ids_list[-1].max()) + 1
2094
+
2095
+ eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
2096
+ llm_pos_ids_list.append(eos_pos)
2097
+ st = offset + 1 + video_len + 1
2098
+ else:
2099
+ audio_bos_pos = np.broadcast_to(np.array([st_idx - 1]), (3, 1))
2100
+ llm_pos_ids_list.append(audio_bos_pos)
2101
+
2102
+ pos_ids, _ = self._compute_interleaved_positions(st_idx, data)
2103
+ llm_pos_ids_list.append(pos_ids)
2104
+ st_idx = int(pos_ids.max()) + 1
2105
+
2106
+ eos_pos = np.broadcast_to(np.array([st_idx]), (3, 1))
2107
+ llm_pos_ids_list.append(eos_pos)
2108
+ llm_pos_ids_list.append(eos_pos)
2109
+
2110
+ video_len = grid_t * grid_h * grid_w
2111
+ audio_len = self._compute_audio_token_count(
2112
+ data["audio_feature_length"]
2113
)
2158
- video_len = video_grid_thw[video_idx].prod() // (spatial_merge_size**2)
2159
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2160
- eos_len = 1
2161
- eos_block = (
2162
- torch.arange(eos_len, dtype=torch.long).view(1, -1).expand(3, -1)
2163
- + st_idx
2164
- )
2165
- llm_pos_ids_list.append(eos_block)
2166
- llm_pos_ids_list.append(eos_block)
2167
- st += text_len + bos_len * 2 + audio_len + video_len + eos_len * 2 # noqa: E501
2168
- audio_idx += 1
2169
- video_idx += 1
2170
- remain_videos -= 1
2171
- remain_audios -= 1
2114
+ st = offset + 2 + video_len + audio_len + 2
2115
2173
- if st < len(input_tokens):
2174
- st_idx = llm_pos_ids_list[-1].max() + 1 if llm_pos_ids_list else 0
2175
- text_len = len(input_tokens) - st
2116
+ if st < seq_len:
2117
+ st_idx = int(llm_pos_ids_list[-1].max()) + 1 if llm_pos_ids_list else 0
2118
+ text_len = seq_len - st
2119
llm_pos_ids_list.append(
2177
- torch.arange(text_len, dtype=torch.long).view(1, -1).expand(3, -1)
2178
- + st_idx
2120
+ np.broadcast_to(np.arange(text_len), (3, text_len)) + st_idx
2121
)
2122
2181
- llm_positions = torch.cat(llm_pos_ids_list, dim=1).reshape(3, -1)
2123
+ llm_positions = np.concatenate(llm_pos_ids_list, axis=1).reshape(3, -1)
2124
if llm_positions.shape[1] != seq_len:
2125
raise RuntimeError("Position ids length mismatch with input ids length")
2126
2185
- mrope_position_delta = llm_positions.max() + 1 - seq_len
2186
- return llm_positions, mrope_position_delta
2127
+ mrope_position_delta = int(llm_positions.max()) + 1 - seq_len
2128
+ return torch.from_numpy(llm_positions), mrope_position_delta
2129
2130
def get_mm_mapping(self) -> MultiModelKeys:
2131
"""