@tduncan / Nvidia-Optimization-Stack-VLLM / commits / 69d4f5ef63

[Bugfix][Multimodal] Fix Qwen3-Omni use_audio_in_video with mixed image/video inputs (#46213)

Signed-off-by: wendadawen <wendadawen@qq.com> Signed-off-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn> Co-authored-by: Tianyu Guo <guoty9@mail2.sysu.edu.cn>

wendadawen committed Jul 17, 2026 at 16:31 UTC 69d4f5ef6323a44621b10e0f269c2775db510cc5
8 files changed +266 -83
tests/models/multimodal/processing/test_qwen2_5_omni_embed.py
+87 -7
@@ -19,6 +19,7 @@ from vllm.model_executor.models.qwen2_5_omni_thinker import (
19 check_interleaved_audio_video,
20 merge_interleaved_embeddings,
21 )
22 +from vllm.multimodal.utils import set_mm_embedding_modality
23
24 # Fake token IDs
25 AUDIO_TOKEN_ID = 1001
@@ -27,6 +28,10 @@ VIDEO_TOKEN_ID = 1003
28 TEXT_TOKEN_ID = 0
29
30
31 +def _mm_embed(shape: tuple[int, ...], value: float, modality: str) -> torch.Tensor:
32 + return set_mm_embedding_modality(torch.full(shape, value), modality)
33 +
34 +
35 # ---------------------------------------------------------------------------
36 # Helpers
37 # ---------------------------------------------------------------------------
@@ -116,6 +121,24 @@ class TestCheckInterleavedAudioVideo:
121 is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
122 )
123
124 + def test_multi_video_with_boundary_tokens(self):
125 + """Two interleaved videos separated by boundary tokens → still True.
126 +
127 + use_audio_in_video expands each video into a local V/A span bounded by
128 + non-pad tokens. A global density check would fail across those gaps.
129 + """
130 + # [text][V A V A][text boundary][V A V A]
131 + first_ids, _ = make_interleaved_seq([2, 2], [2, 2], text_prefix=1)
132 + second_ids, _ = make_interleaved_seq([2, 2], [2, 2], text_prefix=0)
133 + boundary = torch.tensor([TEXT_TOKEN_ID, TEXT_TOKEN_ID])
134 + input_ids = torch.cat([first_ids, boundary, second_ids])
135 + is_multimodal = (input_ids == VIDEO_TOKEN_ID) | (input_ids == AUDIO_TOKEN_ID)
136 + is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
137 + is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
138 + assert check_interleaved_audio_video(
139 + is_video, is_audio, is_video.sum().item(), is_audio.sum().item()
140 + )
141 +
142 def test_batched_non_interleaved_no_false_positive(self):
143 """
144 Regression test for https://github.com/vllm-project/vllm/issues/35394.
@@ -323,13 +346,16 @@ class TestEmbedInputIds:
346 # mm_embeds come in [video, audio] order (video feature first in
347 # mm_features when positions are the same for use_audio_in_video)
348 mm_embeds = [
326 - torch.full((video_n, hidden), video_val),
327 - torch.full((audio_n, hidden), audio_val),
349 + _mm_embed((video_n, hidden), video_val, "video"),
350 + _mm_embed((audio_n, hidden), audio_val, "audio"),
351 ]
352
353 model, _ = make_mock_model(hidden)
354 + # Modalities are attached on the embedding tensors (as in encoder gather).
355 result = model.embed_input_ids(
332 - input_ids, mm_embeds, is_multimodal=is_multimodal
356 + input_ids,
357 + mm_embeds,
358 + is_multimodal=is_multimodal,
359 )
360
361 video_pos = (input_ids == VIDEO_TOKEN_ID).nonzero(as_tuple=True)[0]
@@ -362,8 +388,8 @@ class TestMergeInterleavedEmbeddings:
388
389 inputs_embeds = torch.zeros(len(input_ids), hidden)
390 mm_embeds = [
365 - torch.full((num_video, hidden), 30.0),
366 - torch.full((num_audio, hidden), 10.0),
391 + _mm_embed((num_video, hidden), 30.0, "video"),
392 + _mm_embed((num_audio, hidden), 10.0, "audio"),
393 ]
394
395 result = merge_interleaved_embeddings(
@@ -372,8 +398,6 @@ class TestMergeInterleavedEmbeddings:
398 is_video,
399 is_audio,
400 is_multimodal,
375 - num_video,
376 - num_audio,
401 )
402
403 video_pos = is_video.nonzero(as_tuple=True)[0]
@@ -381,6 +405,62 @@ class TestMergeInterleavedEmbeddings:
405 assert result[video_pos].allclose(torch.full((num_video, hidden), 30.0))
406 assert result[audio_pos].allclose(torch.full((num_audio, hidden), 10.0))
407
408 + def test_image_and_video_mixed(self):
409 + """Image embeddings must not be misclassified as video."""
410 + hidden = 4
411 + # [text][I I][V A V A]
412 + tokens = (
413 + [TEXT_TOKEN_ID] * 2
414 + + [IMAGE_TOKEN_ID] * 2
415 + + [VIDEO_TOKEN_ID, AUDIO_TOKEN_ID, VIDEO_TOKEN_ID, AUDIO_TOKEN_ID]
416 + )
417 + input_ids = torch.tensor(tokens)
418 + is_multimodal = (
419 + (input_ids == IMAGE_TOKEN_ID)
420 + | (input_ids == VIDEO_TOKEN_ID)
421 + | (input_ids == AUDIO_TOKEN_ID)
422 + )
423 + is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
424 + is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
425 + is_image = is_multimodal & (input_ids == IMAGE_TOKEN_ID)
426 +
427 + inputs_embeds = torch.zeros(len(input_ids), hidden)
428 + mm_embeds = [
429 + _mm_embed((2, hidden), 20.0, "image"),
430 + _mm_embed((2, hidden), 30.0, "video"),
431 + _mm_embed((2, hidden), 10.0, "audio"),
432 + ]
433 + result = merge_interleaved_embeddings(
434 + inputs_embeds,
435 + mm_embeds,
436 + is_video,
437 + is_audio,
438 + is_multimodal,
439 + )
440 + assert result[is_image.nonzero(as_tuple=True)[0]].allclose(
441 + torch.full((2, hidden), 20.0)
442 + )
443 + assert result[is_video.nonzero(as_tuple=True)[0]].allclose(
444 + torch.full((2, hidden), 30.0)
445 + )
446 + assert result[is_audio.nonzero(as_tuple=True)[0]].allclose(
447 + torch.full((2, hidden), 10.0)
448 + )
449 +
450 + def test_missing_modality_raises(self):
451 + hidden = 2
452 + input_ids, is_multimodal = make_interleaved_seq([2], [2])
453 + is_video = is_multimodal & (input_ids == VIDEO_TOKEN_ID)
454 + is_audio = is_multimodal & (input_ids == AUDIO_TOKEN_ID)
455 + with pytest.raises(ValueError, match="Missing modality"):
456 + merge_interleaved_embeddings(
457 + torch.zeros(len(input_ids), hidden),
458 + [torch.zeros(2, hidden), torch.zeros(2, hidden)],
459 + is_video,
460 + is_audio,
461 + is_multimodal,
462 + )
463 +
464
465 if __name__ == "__main__":
466 pytest.main([__file__, "-v"])
tests/v1/worker/test_encoder_runner.py
+26
@@ -78,6 +78,7 @@ def test_draft_lookahead_uses_boundary_feature_when_cached():
78
79 # f0 covers positions 0..6 (+1 skew); f1's first embed covers position 7.
80 assert len(mm_embeds) == 2
81 + assert [e.modality for e in mm_embeds] == ["image", "image"]
82 assert bool(is_mm_embed[7])
83 assert int(is_mm_embed.sum()) == 8
84
@@ -94,6 +95,7 @@ def test_draft_lookahead_tolerates_missing_boundary_feature():
95
96 # Only f0 is gathered; f1's boundary position falls back silently.
97 assert len(mm_embeds) == 1
98 + assert [e.modality for e in mm_embeds] == ["image"]
99 assert not bool(is_mm_embed[7])
100 assert int(is_mm_embed.sum()) == 7
101
@@ -154,4 +156,28 @@ def test_multi_request_batch_gathers_per_request(draft_lookahead):
156 # Both requests contribute a feature; with the +1 skew each marks 7 of its
157 # 8 positions (the skew drops one), otherwise all 8.
158 assert len(mm_embeds) == 2
159 + assert [e.modality for e in mm_embeds] == ["image", "image"]
160 assert int(is_mm_embed.sum()) == (14 if draft_lookahead else 16)
161 +
162 +
163 +def test_gather_preserves_mixed_modalities():
164 + """Modalities must be attached on tensors in gather order."""
165 + video = MultiModalFeatureSpec(
166 + data=None,
167 + modality="video",
168 + identifier="v0",
169 + mm_position=PlaceholderRange(offset=0, length=4),
170 + )
171 + audio = MultiModalFeatureSpec(
172 + data=None,
173 + modality="audio",
174 + identifier="a0",
175 + mm_position=PlaceholderRange(offset=4, length=4),
176 + )
177 + runner = _make_runner([video, audio], cached=[video, audio])
178 +
179 + mm_embeds, is_mm_embed = _gather(runner, num_scheduled=8, draft_lookahead=0)
180 +
181 + assert len(mm_embeds) == 2
182 + assert [e.modality for e in mm_embeds] == ["video", "audio"]
183 + assert int(is_mm_embed.sum()) == 8
vllm/model_executor/models/qwen2_5_omni_thinker.py
+77 -66
@@ -138,42 +138,44 @@ def check_interleaved_audio_video(
138 num_audio: int,
139 ) -> bool:
140 """
141 - Check if video and audio positions are interleaved in the multimodal region.
142 -
143 - Returns True only for the use_audio_in_video=True case, where video and
144 - audio tokens alternate within a single contiguous region with no gaps.
145 -
146 - A simple range-overlap check produces false positives when multiple
147 - non-interleaved requests are batched together: audio tokens from request N
148 - fall between video tokens from request N and request N+1, making the
149 - global ranges overlap even though each individual request is non-interleaved.
150 -
151 - To distinguish true interleaving from this batching artefact we require
152 - that every position in the combined [first_VA, last_VA] range is occupied
153 - by either a video or an audio token (no text/image gaps).
141 + Check if video and audio positions are interleaved in any per-video span.
142 +
143 + For use_audio_in_video=True, each video placeholder is expanded into one
144 + local span containing only video/audio pad tokens, bounded by non-pad
145 + tokens such as audio_start/audio_end. Check each contiguous V/A span
146 + independently instead of requiring all V/A tokens in the whole sequence to
147 + form one global dense range; otherwise multi-video requests can be
148 + misclassified as non-interleaved because of the boundary tokens between
149 + videos.
150 """
151 if num_video == 0 or num_audio == 0:
152 return False
153
158 - video_pos = is_video.nonzero(as_tuple=True)[0]
159 - audio_pos = is_audio.nonzero(as_tuple=True)[0]
160 -
161 - # Quick range-overlap pre-check (necessary but not sufficient).
162 - if not (
163 - video_pos[0].item() < audio_pos[-1].item()
164 - and audio_pos[0].item() < video_pos[-1].item()
165 - ):
154 + va_pos = (is_video | is_audio).nonzero(as_tuple=True)[0]
155 + if len(va_pos) == 0:
156 return False
157
168 - # Density check: for true use_audio_in_video interleaving every position
169 - # in the combined span is a video or audio token. Batched non-interleaved
170 - # requests have text/image tokens between the per-request V and A blocks.
171 - # combined_start/end encompass all V/A tokens, so num_video + num_audio
172 - # equals the number of V/A tokens in range; compare directly to span size.
173 - combined_start = min(video_pos[0].item(), audio_pos[0].item())
174 - combined_end = max(video_pos[-1].item(), audio_pos[-1].item())
175 - total_in_range = combined_end - combined_start + 1
176 - return (num_video + num_audio) == total_in_range
158 + span_start = 0
159 + for span_end in range(1, len(va_pos) + 1):
160 + is_last = span_end == len(va_pos)
161 + if not is_last and va_pos[span_end].item() == va_pos[span_end - 1].item() + 1:
162 + continue
163 +
164 + span = va_pos[span_start:span_end]
165 + span_is_video = is_video[span]
166 + span_is_audio = is_audio[span]
167 + if span_is_video.any() and span_is_audio.any():
168 + video_offsets = span_is_video.nonzero(as_tuple=True)[0]
169 + audio_offsets = span_is_audio.nonzero(as_tuple=True)[0]
170 + if (
171 + video_offsets[0].item() < audio_offsets[-1].item()
172 + and audio_offsets[0].item() < video_offsets[-1].item()
173 + ):
174 + return True
175 +
176 + span_start = span_end
177 +
178 + return False
179
180
181 def merge_interleaved_embeddings(
@@ -182,58 +184,69 @@ def merge_interleaved_embeddings(
184 is_video: torch.Tensor,
185 is_audio: torch.Tensor,
186 is_multimodal: torch.Tensor,
185 - num_video: int,
186 - num_audio: int,
187 ) -> torch.Tensor:
188 """
189 Merge embeddings for interleaved audio-in-video sequences.
190
191 When use_audio_in_video=True, video and audio tokens are interleaved in
192 the token sequence, but embeddings are provided as separate contiguous
193 - tensors (video first, then audio). This function reorders video and audio
194 - embeddings to match sequence position order and scatters them efficiently.
193 + tensors. This function scatters each modality by the ``modality`` attribute
194 + attached to each embedding tensor (set during encoder gather) and also
195 + supports image embeddings in the same interleaved request.
196
197 Args:
198 inputs_embeds: The input embeddings tensor to merge into.
198 - multimodal_embeddings: List of embedding tensors (video, audio, other).
199 + multimodal_embeddings: List of embedding tensors (video, audio, image).
200 is_video: Boolean mask for video token positions.
201 is_audio: Boolean mask for audio token positions.
202 is_multimodal: Boolean mask for all multimodal token positions.
202 - num_video: Total count of video tokens.
203 - num_audio: Total count of audio tokens.
203
204 Returns:
205 The merged inputs_embeds tensor with multimodal embeddings scattered
206 to their correct positions.
207 """
209 - # Categorize embeddings by modality based on token counts.
210 - # Embeddings come grouped by modality but order varies (e.g., image, video, audio
211 - # or video, audio depending on input kwargs order).
212 - video_embeds: list[torch.Tensor] = []
213 - audio_embeds: list[torch.Tensor] = []
214 - other_embeds: list[torch.Tensor] = []
215 - video_remaining = num_video
216 - audio_remaining = num_audio
217 -
218 - for emb in multimodal_embeddings:
219 - n = emb.shape[0]
220 - if video_remaining > 0 and n <= video_remaining:
221 - video_embeds.append(emb)
222 - video_remaining -= n
223 - elif audio_remaining > 0 and n <= audio_remaining:
224 - audio_embeds.append(emb)
225 - audio_remaining -= n
226 - else:
227 - other_embeds.append(emb)
208 + from vllm.multimodal.utils import get_mm_embedding_modalities
209 +
210 + def _merge_embedding_group(
211 + mask: torch.Tensor,
212 + embeddings: Sequence[torch.Tensor],
213 + modality: str,
214 + ) -> None:
215 + num_expected_tokens = mask.sum().item()
216 + num_actual_tokens = sum(embedding.shape[0] for embedding in embeddings)
217 + if num_actual_tokens != num_expected_tokens:
218 + raise ValueError(
219 + f"Attempted to assign {num_actual_tokens} {modality} tokens "
220 + f"to {num_expected_tokens} placeholders"
221 + )
222 + if embeddings:
223 + inputs_embeds[mask] = torch.cat(list(embeddings), dim=0)
224 +
225 + # Qwen Omni multimodal placeholders are image/video/audio; after excluding
226 + # video and audio positions, the remaining multimodal positions are images.
227 + is_image = is_multimodal & ~is_video & ~is_audio
228 +
229 + embedding_modalities = get_mm_embedding_modalities(multimodal_embeddings)
230 +
231 + video_embeds = [
232 + embedding
233 + for embedding, modality in zip(multimodal_embeddings, embedding_modalities)
234 + if modality == "video"
235 + ]
236 + audio_embeds = [
237 + embedding
238 + for embedding, modality in zip(multimodal_embeddings, embedding_modalities)
239 + if modality == "audio"
240 + ]
241 + image_embeds = [
242 + embedding
243 + for embedding, modality in zip(multimodal_embeddings, embedding_modalities)
244 + if modality == "image"
245 + ]
246
229 - # Scatter each modality to its positions
230 - if video_embeds:
231 - inputs_embeds[is_video] = torch.cat(video_embeds, dim=0)
232 - if audio_embeds:
233 - inputs_embeds[is_audio] = torch.cat(audio_embeds, dim=0)
234 - if other_embeds:
235 - other_mask = is_multimodal & ~is_video & ~is_audio
236 - inputs_embeds[other_mask] = torch.cat(other_embeds, dim=0)
247 + _merge_embedding_group(is_video, video_embeds, "video")
248 + _merge_embedding_group(is_audio, audio_embeds, "audio")
249 + _merge_embedding_group(is_image, image_embeds, "image")
250
251 return inputs_embeds
252
@@ -1497,8 +1510,6 @@ class Qwen2_5OmniThinkerForConditionalGeneration(
1510 is_video,
1511 is_audio,
1512 is_multimodal,
1500 - num_video,
1501 - num_audio,
1513 )
1514
1515 # Default: standard merge (no interleaving), same as parent class
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+5 -5
@@ -1833,9 +1833,11 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1833 # both the deepstack path and the final embedding merge.
1834 video_token_id = self.config.video_token_id
1835 audio_token_id = self.config.audio_token_id
1836 + image_token_id = self.config.image_token_id
1837 input_ids_cpu = input_ids.cpu()
1838 is_video = is_multimodal & (input_ids_cpu == video_token_id)
1839 is_audio = is_multimodal & (input_ids_cpu == audio_token_id)
1840 + is_image = is_multimodal & (input_ids_cpu == image_token_id)
1841 num_video = is_video.sum().item()
1842 num_audio = is_audio.sum().item()
1843
@@ -1856,9 +1858,9 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1858 multimodal_embeddings_multiscale = []
1859
1860 if is_interleaved:
1859 - # Use input_ids-based mask for correct vision positions
1860 - # when audio and video tokens are interleaved.
1861 - is_vision = is_video.clone()
1861 + # Use input_ids-based mask for all vision positions when
1862 + # audio and video tokens are interleaved.
1863 + is_vision = is_video | is_image
1864 else:
1865 is_vision = torch.zeros_like(is_multimodal)
1866 mm_positions = torch.nonzero(is_multimodal, as_tuple=True)[0]
@@ -1917,8 +1919,6 @@ class Qwen3OmniMoeThinkerForConditionalGeneration(
1919 is_video,
1920 is_audio,
1921 is_multimodal,
1920 - num_video,
1921 - num_audio,
1922 )
1923
1924 # Default: standard merge (no interleaving), same as parent class.
vllm/multimodal/utils.py
+38
@@ -346,3 +346,41 @@ def fetch_video(
346 allowed_local_media_path="/",
347 )
348 return media_connector.fetch_video(video_url)
349 +
350 +
351 +def set_mm_embedding_modality(embed: "torch.Tensor", modality: str) -> "torch.Tensor":
352 + """Attach modality metadata to a gathered multimodal embedding tensor.
353 +
354 + Used by interleaved Omni merge paths that need to group embeddings by
355 + modality without threading a parallel modalities list through
356 + ``embed_input_ids``.
357 + """
358 + embed.modality = modality # type: ignore[attr-defined]
359 + return embed
360 +
361 +
362 +def copy_mm_embedding_modality(
363 + src: "torch.Tensor", dst: "torch.Tensor"
364 +) -> "torch.Tensor":
365 + """Copy ``modality`` from ``src`` onto ``dst`` if present."""
366 + modality = getattr(src, "modality", None)
367 + if modality is not None:
368 + dst.modality = modality # type: ignore[attr-defined]
369 + return dst
370 +
371 +
372 +def get_mm_embedding_modalities(
373 + multimodal_embeddings: Sequence["torch.Tensor"],
374 +) -> list[str]:
375 + """Collect per-embedding modalities previously set on the tensors."""
376 + modalities: list[str] = []
377 + for i, emb in enumerate(multimodal_embeddings):
378 + modality = getattr(emb, "modality", None)
379 + if modality is None:
380 + raise ValueError(
381 + f"Missing modality on multimodal embedding at index {i}. "
382 + "Encoder gather must set embed.modality before interleaved "
383 + "audio-in-video merge."
384 + )
385 + modalities.append(modality)
386 + return modalities
vllm/v1/worker/gpu/mm/encoder_runner.py
+8 -1
@@ -5,7 +5,11 @@ import torch
5
6 from vllm.model_executor.models.interfaces import SupportsMultiModal, supports_realtime
7 from vllm.multimodal.inputs import MultiModalKwargsItem
8 -from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs
8 +from vllm.multimodal.utils import (
9 + get_mm_features_in_window,
10 + group_and_batch_mm_kwargs,
11 + set_mm_embedding_modality,
12 +)
13 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
14 from vllm.v1.worker.utils import sanity_check_mm_encoder_outputs
15
@@ -136,6 +140,9 @@ class EncoderRunner:
140 else:
141 mm_embeds_item = encoder_output[start_idx:end_idx]
142
143 + # Attach modality for Omni interleaved merge (collected on demand).
144 + set_mm_embedding_modality(mm_embeds_item, mm_feature.modality)
145 +
146 req_start_pos = query_start_loc[i] + start_pos - cur_query_start
147 is_mm_embed[req_start_pos + start_idx : req_start_pos + end_idx] |= (
148 True if is_embed is None else is_embed
vllm/v1/worker/gpu/model_states/mm_pruning.py
+13 -3
@@ -5,7 +5,10 @@ import torch.nn as nn
5
6 from vllm.config import ModelConfig
7 from vllm.model_executor.models.interfaces import supports_multimodal_pruning
8 -from vllm.multimodal.utils import get_mm_features_in_window
8 +from vllm.multimodal.utils import (
9 + copy_mm_embedding_modality,
10 + get_mm_features_in_window,
11 +)
12 from vllm.v1.worker.gpu.input_batch import InputBatch
13 from vllm.v1.worker.gpu.mm.encoder_cache import EncoderCache
14 from vllm.v1.worker.gpu.mm.rope import RopeState
@@ -42,7 +45,11 @@ class MultiModalPruner:
45 speculator reuses the target's already-recomputed positions, hence there is
46 no position write-back here.
47 """
45 - return [mm[:, : self.inputs_embeds_size] for mm in mm_embeds]
48 + stripped: list[torch.Tensor] = []
49 + for mm in mm_embeds:
50 + out = mm[:, : self.inputs_embeds_size]
51 + stripped.append(copy_mm_embedding_modality(mm, out))
52 + return stripped
53
54 def recompute(
55 self,
@@ -82,7 +89,10 @@ class MultiModalPruner:
89 num_computed_tokens=num_computed,
90 )
91 self.rope_state.update_prefill_positions(req_idx, new_positions, delta)
85 - cleaned.extend(req_cleaned)
92 + cleaned.extend(
93 + copy_mm_embedding_modality(src, dst)
94 + for src, dst in zip(req_embeds, req_cleaned)
95 + )
96
97 assert pos == len(mm_embeds)
98 return cleaned
vllm/v1/worker/gpu_model_runner.py
+12 -1
@@ -105,7 +105,12 @@ from vllm.multimodal.inputs import (
105 MultiModalKwargsItem,
106 PlaceholderRange,
107 )
108 -from vllm.multimodal.utils import get_mm_features_in_window, group_and_batch_mm_kwargs
108 +from vllm.multimodal.utils import (
109 + copy_mm_embedding_modality,
110 + get_mm_features_in_window,
111 + group_and_batch_mm_kwargs,
112 + set_mm_embedding_modality,
113 +)
114 from vllm.platforms import current_platform
115 from vllm.pooling_params import PoolingParams
116 from vllm.sampling_params import SamplingType
@@ -3245,11 +3250,13 @@ class GPUModelRunner(
3250 is_mm_embed[
3251 req_start_pos + start_idx : req_start_pos + end_idx
3252 ] |= is_embed
3253 + set_mm_embedding_modality(mm_embeds_item, mm_feature.modality)
3254 mm_embeds_req.append(mm_embeds_item)
3255
3256 if self.is_multimodal_pruning_enabled and self.uses_mrope:
3257 assert req_state.mrope_positions is not None
3258 should_sync_mrope_positions = True
3259 + old_mm_embeds_req = mm_embeds_req
3260 mm_embeds_req, new_mrope_positions, new_delta = (
3261 self.model.recompute_mrope_positions(
3262 input_ids=req_state.prompt_token_ids,
@@ -3258,6 +3265,10 @@ class GPUModelRunner(
3265 num_computed_tokens=req_state.num_computed_tokens,
3266 )
3267 )
3268 + mm_embeds_req = [
3269 + copy_mm_embedding_modality(src, dst)
3270 + for src, dst in zip(old_mm_embeds_req, mm_embeds_req)
3271 + ]
3272 req_state.mrope_positions.copy_(new_mrope_positions)
3273 req_state.mrope_position_delta = new_delta
3274