[EPD] Support metadata-only audio inputs (#57887)
Signed-off-by: Tianyu Guo <guoty@inferact.ai> Co-authored-by: Cyrus Leung <tlleungac@connect.ust.hk>
Tianyu Guo committed
Sep 23, 2026 at 09:03 UTC
3ee2906f47afdfc1d792e04995c31eb6ecf2bdef
14 files changed
+387
-27
docs/features/disagg_encoder.md
+13
@@ -61,6 +61,19 @@ and compares the disaggregated outputs against it. Set `MODEL` to use another mo
61
or `MOONCAKE_EC_PROTOCOL=rdma` to use RDMA on supported hardware. See the script for
62
additional configuration options.
63
64
+### Audio inputs
65
+
66
+The Python EPD proxy also rewrites audio in `/v1/chat/completions` into
67
+metadata-only references for Qwen2-Audio, AudioFlamingo3, Ultravox,
68
+Qwen2.5-Omni, and Qwen3-Omni. The encoder publishes `audio_num_tokens`
69
+(the placeholder token count of each audio) together with any feature
70
+lengths the model needs; the consumer uses them to reconstruct the audio
71
+placeholders and receives the embeddings through its EC connector without
72
+repeating audio preprocessing.
73
+
74
+This covers pure audio inputs, not video with an embedded audio track. It does
75
+not add `/v1/audio/transcriptions` or realtime routes to the example proxy.
76
+
77
---
78
79
## 3 Test Script
examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py
+34
-11
@@ -480,6 +480,17 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
480
the grid its processor actually produced), so the grid is never re-derived
481
here -- a second derivation could disagree with the encoder's.
482
"""
483
+ mm_items = extract_mm_items(req_data)
484
+ # Audio processors may consume placeholders from other modalities too.
485
+ # Keep the whole request raw if any item cannot be rewritten safely.
486
+ control_fields = {"mm_hash", "ec_mm_hash", "transfer_id"}
487
+ raw_audio_request = any(
488
+ item["type"] in {"audio_url", "input_audio"} for item in mm_items
489
+ ) and any(
490
+ not item_meta.get(i, {}).get("mm_hash")
491
+ or not (item_meta.get(i, {}).keys() - control_fields)
492
+ for i in range(len(mm_items))
493
+ )
494
rewritten = 0
495
transfer_items = []
496
idx = 0
@@ -499,6 +510,18 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
510
item_uuid = meta.pop("mm_hash", None)
511
ec_mm_hash = meta.pop("ec_mm_hash", None) or item_uuid
512
transfer_id = meta.pop("transfer_id", None)
513
+ if raw_audio_request:
514
+ new_content.append(
515
+ {
516
+ **item,
517
+ "uuid": item_uuid or item.get("uuid") or content_uuid(item),
518
+ }
519
+ )
520
+ if transfer_id is not None:
521
+ transfer_items.append(
522
+ {"mm_hash": ec_mm_hash, "transfer_id": transfer_id}
523
+ )
524
+ continue
525
# Whatever keys the encoder reported are the metadata its model
526
# declared as needed to size the placeholder range; the proxy does
527
# not need to know their names.
@@ -532,9 +555,10 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict:
555
rewritten += 1
556
new_messages.append({**msg, "content": new_content})
557
535
- if not rewritten:
558
+ if not rewritten and not raw_audio_request:
559
return req_data
537
- logger.info("Rewrote %d media item(s) as metadata references", rewritten)
560
+ if rewritten:
561
+ logger.info("Rewrote %d media item(s) as metadata references", rewritten)
562
rewritten_request = {**req_data, "messages": new_messages}
563
if transfer_items:
564
ec_transfer_params = dict(req_data.get("ec_transfer_params") or {})
@@ -632,7 +656,7 @@ async def fanout_encoder_primer(
656
content = []
657
for idx in indices:
658
item = mm_items[idx]
635
- item_uuid = None if NO_REWRITE else content_uuid(item)
659
+ item_uuid = None if NO_REWRITE else (item.get("uuid") or content_uuid(item))
660
if item_uuid is not None:
661
item_uuids[idx] = item_uuid
662
item_transfer_ids[idx] = uuid.uuid4().hex
@@ -756,14 +780,13 @@ async def fanout_encoder_primer(
780
raise HTTPException(502, "Encoder metadata cannot be matched")
781
ec_mm_hash, reported = matched
782
metadata = reported.get("metadata") or {}
759
- if metadata:
760
- item_meta[idx] = {
761
- **metadata,
762
- "mm_hash": item_uuids.get(idx, ec_mm_hash),
763
- "ec_mm_hash": ec_mm_hash,
764
- }
765
- if idx in item_transfer_ids:
766
- item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
783
+ item_meta[idx] = {
784
+ **metadata,
785
+ "mm_hash": item_uuids.get(idx, ec_mm_hash),
786
+ "ec_mm_hash": ec_mm_hash,
787
+ }
788
+ if idx in item_transfer_ids:
789
+ item_meta[idx]["transfer_id"] = item_transfer_ids[idx]
790
# Whatever the encoder reported alongside `metadata` is the
791
# connector's own handle on the published embedding (for NIXL,
792
# peer_host/peer_port/size_bytes). The decoder's connector
tests/model_executor/test_qwen3_omni.py
+5
@@ -144,6 +144,11 @@ def test_qwen3_omni_get_updates_use_audio_in_video(
144
# Create a mock context
145
mock_ctx = Mock(spec=InputProcessingContext)
146
mock_ctx.tokenizer = mock_tokenizer
147
+ # `model_config` is an instance attribute, so it is not covered by the
148
+ # spec; the data parser reads it for `allow_missing_mm_embeddings`.
149
+ mock_ctx.model_config = SimpleNamespace(
150
+ multimodal_config=SimpleNamespace(allow_missing_mm_embeddings=False)
151
+ )
152
153
# Create processing info
154
info = Qwen3OmniMoeThinkerProcessingInfo(mock_ctx)
tests/models/multimodal/processing/test_common.py
+82
-1
@@ -1,13 +1,16 @@
1
# SPDX-License-Identifier: Apache-2.0
2
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
3
4
+import json
5
from collections.abc import Set as AbstractSet
6
from functools import partial
7
8
import numpy as np
9
import pytest
10
+import torch
11
from PIL import Image
12
13
+from tests.models.utils import build_model_context
14
from vllm.config import ModelConfig
15
from vllm.config.multimodal import (
16
AudioDummyOptions,
@@ -16,10 +19,15 @@ from vllm.config.multimodal import (
19
MultiModalDummyOptions,
20
VideoDummyOptions,
21
)
22
+from vllm.distributed.ec_transfer.ec_connector.utils import (
23
+ PlaceholderMetadataResolver,
24
+ collect_ec_item_metadata,
25
+)
26
+from vllm.entrypoints.chat_utils import _get_embeds_data
27
from vllm.inputs import MultiModalDataDict, MultiModalInput
28
from vllm.multimodal import MULTIMODAL_REGISTRY
29
from vllm.multimodal.cache import MultiModalProcessorOnlyCache
22
-from vllm.multimodal.inputs import batched_tensors_equal
30
+from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal
31
from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext
32
from vllm.platforms import current_platform
33
from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config
@@ -193,6 +201,79 @@ def get_token_prompt(
201
return inputs.prompt
202
203
204
+@pytest.mark.parametrize(
205
+ "model_id,durations",
206
+ [
207
+ ("Qwen/Qwen2-Audio-7B-Instruct", [1.0, 2.3]),
208
+ ("nvidia/audio-flamingo-3-hf", [31.0]),
209
+ ("fixie-ai/ultravox-v0_5-llama-3_2-1b", [1.0, 31.0]),
210
+ ("Qwen/Qwen2.5-Omni-3B", [1.0, 2.3]),
211
+ ("Qwen/Qwen3-Omni-30B-A3B-Instruct", [1.0, 2.3]),
212
+ ],
213
+)
214
+def test_audio_metadata_only_roundtrip(model_id, durations, monkeypatch):
215
+ """EC metadata preserves the raw-audio prompt without running the HF processor."""
216
+ ctx = build_model_context(model_id, limit_mm_per_prompt={"audio": len(durations)})
217
+ mm_config = ctx.model_config.multimodal_config
218
+ mm_config.enable_mm_embeds = True
219
+ mm_config.allow_missing_mm_embeddings = True
220
+ processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
221
+ sample_rate = processor.info.get_feature_extractor().sampling_rate
222
+ audios = [
223
+ np.zeros(int(duration * sample_rate), dtype=np.float32)
224
+ for duration in durations
225
+ ]
226
+ prompt = processor.dummy_inputs.get_dummy_text({"audio": len(audios)})
227
+ raw = processor(
228
+ prompt,
229
+ mm_items=processor.info.parse_mm_data({"audio": audios}),
230
+ hf_processor_mm_kwargs={},
231
+ )
232
+ features = [
233
+ MultiModalFeatureSpec(
234
+ data=item,
235
+ modality="audio",
236
+ identifier=str(index),
237
+ mm_position=position,
238
+ )
239
+ for index, (item, position) in enumerate(
240
+ zip(raw["mm_kwargs"]["audio"], raw["mm_placeholders"]["audio"])
241
+ )
242
+ ]
243
+ published = collect_ec_item_metadata(
244
+ features, PlaceholderMetadataResolver(ctx.model_config)
245
+ )
246
+ metadata = json.loads(
247
+ json.dumps([entry["metadata"] for entry in published.values()])
248
+ )
249
+ assert len(metadata) == len(audios)
250
+ assert all(metadata)
251
+ audio_data = _get_embeds_data("audio", metadata, processor)
252
+
253
+ def unexpected_preprocessing(*args, **kwargs):
254
+ pytest.fail("The EC consumer must not preprocess audio")
255
+
256
+ monkeypatch.setattr(processor, "_call_hf_processor", unexpected_preprocessing)
257
+ remote = processor(
258
+ prompt,
259
+ mm_items=processor.info.parse_mm_data({"audio": audio_data}),
260
+ hf_processor_mm_kwargs={},
261
+ )
262
+ assert remote["prompt_token_ids"] == raw["prompt_token_ids"]
263
+ assert remote["mm_placeholders"] == raw["mm_placeholders"]
264
+ # Omni also consumes these lengths when constructing M-RoPE positions.
265
+ if "audio_feature_lengths" in audio_data:
266
+ assert torch.equal(
267
+ remote["mm_kwargs"].get_data()["audio_feature_lengths"],
268
+ raw["mm_kwargs"].get_data()["audio_feature_lengths"],
269
+ )
270
+
271
+ mm_config.allow_missing_mm_embeddings = False
272
+ ordinary_processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config)
273
+ with pytest.raises(ValueError, match="should contain the fields"):
274
+ ordinary_processor.info.parse_mm_data({"audio": audio_data})
275
+
276
+
277
def random_vision_chunk(
278
rng: np.random.RandomState,
279
min_wh: int,
tests/multimodal/test_parse.py
+39
@@ -15,6 +15,45 @@ from vllm.multimodal.parse import (
15
H, W = 480, 640
16
17
18
+class AudioMetadataParser(MultiModalDataParser):
19
+ embedding_fields = {
20
+ "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
21
+ }
22
+
23
+
24
+@pytest.mark.parametrize("allow_missing", [False, True])
25
+def test_audio_metadata_requires_ec_consumer(allow_missing):
26
+ parser = AudioMetadataParser(allow_missing_mm_embeddings=allow_missing)
27
+ data = {"audio_num_tokens": torch.tensor([[3], [5]])}
28
+ if not allow_missing:
29
+ with pytest.raises(ValueError, match="audio_embeds"):
30
+ parser.parse_mm_data({"audio": data})
31
+ else:
32
+ items = parser.parse_mm_data({"audio": data})["audio"]
33
+ assert len(items) == 2
34
+ assert items.get(1)["audio_num_tokens"].item() == 5
35
+ assert items.get_processor_data() == {}
36
+
37
+
38
+@pytest.mark.parametrize("counts", [[0], [-1], [1.5], [True], [[1, 2]]])
39
+def test_audio_metadata_rejects_invalid_token_counts(counts):
40
+ parser = AudioMetadataParser(allow_missing_mm_embeddings=True)
41
+ with pytest.raises(ValueError, match="positive integer"):
42
+ parser.parse_mm_data({"audio": {"audio_num_tokens": torch.tensor(counts)}})
43
+
44
+
45
+def test_audio_metadata_checks_supplied_embedding_lengths():
46
+ parser = AudioMetadataParser()
47
+ data = {
48
+ "audio_num_tokens": torch.tensor([3, 5]),
49
+ "audio_embeds": [torch.zeros(3, 8), torch.zeros(5, 8)],
50
+ }
51
+ assert len(parser.parse_mm_data({"audio": data})["audio"]) == 2
52
+ data["audio_num_tokens"] = torch.tensor([3, 4])
53
+ with pytest.raises(ValueError, match="does not match"):
54
+ parser.parse_mm_data({"audio": data})
55
+
56
+
57
@pytest.mark.parametrize(
58
"image",
59
[
tests/v1/ec_connector/unit/test_ec_example_connector.py
+16
@@ -56,6 +56,22 @@ def test_metadata_preserves_repeated_item_positions():
56
}
57
58
59
+def test_audio_metadata_survives_processor_cache_hit():
60
+ resolver = Mock()
61
+ resolver.fields_for.return_value = {"audio_num_tokens"}
62
+ feature = MultiModalFeatureSpec(
63
+ data=None,
64
+ modality="audio",
65
+ identifier="audio",
66
+ mm_position=PlaceholderRange(
67
+ offset=0, length=4, is_embed=torch.tensor([False, True, True, False])
68
+ ),
69
+ )
70
+ assert collect_ec_item_metadata([feature], resolver) == {
71
+ "audio": {"metadata": {"audio_num_tokens": [2]}, "item_indices": [0]}
72
+ }
73
+
74
+
75
@pytest.fixture
76
def temp_storage(tmp_path):
77
"""Fixture providing temporary storage path."""
tests/v1/ec_connector/unit/test_epd_proxy_retry.py
+81
-1
@@ -87,6 +87,82 @@ class _EncoderSession:
87
return _EncoderResponse(self._replies.pop(0))
88
89
90
+@pytest.mark.parametrize("second_type", ["image_url", "audio_url", "input_audio"])
91
+@pytest.mark.parametrize("transfer", ["example", "nixl", "mooncake"])
92
+def test_partial_audio_metadata_falls_back_with_all_transfers(
93
+ proxy, monkeypatch, second_type, transfer
94
+):
95
+ """Never mix metadata-only audio with raw media, or lose EC identities."""
96
+ import copy
97
+
98
+ seen = []
99
+ reported = {}
100
+
101
+ class Session:
102
+ async def post(self, url, data=None, headers=None):
103
+ body = msgspec.json.decode(data)
104
+ seen.append(body)
105
+ item = body["messages"][0]["content"][0]
106
+ mm_hash = item["uuid"] + "-processed"
107
+ entry: dict[str, Any] = {
108
+ "metadata": {"audio_feature_lengths": [100]} if len(seen) == 1 else {},
109
+ "item_indices": [0],
110
+ }
111
+ if transfer == "nixl":
112
+ entry.update(peer_host="encoder", peer_port=1234, size_bytes=4096)
113
+ reported[mm_hash] = entry
114
+ return _EncoderResponse({mm_hash: entry})
115
+
116
+ second = (
117
+ {"type": second_type, "input_audio": {"data": "large", "format": "wav"}}
118
+ if second_type == "input_audio"
119
+ else {"type": second_type, second_type: {"url": "large"}}
120
+ )
121
+ body = {
122
+ "messages": [
123
+ {
124
+ "role": "user",
125
+ "content": [
126
+ {
127
+ "type": "audio_url",
128
+ "audio_url": {"url": "small"},
129
+ "uuid": "user-audio",
130
+ },
131
+ second,
132
+ ],
133
+ }
134
+ ],
135
+ "mm_processor_kwargs": {"sampling_rate": 16000},
136
+ }
137
+ original = copy.deepcopy(body)
138
+ monkeypatch.setattr(proxy, "encode_session", Session())
139
+ monkeypatch.setattr(proxy, "NO_REWRITE", False)
140
+ consumer = "tcp://consumer:1" if transfer == "mooncake" else None
141
+ prepared, _, _ = asyncio.run(
142
+ proxy.prepare_for_decode(body, "r", ["http://encoder"], "", consumer)
143
+ )
144
+ assert body == original
145
+ assert prepared["mm_processor_kwargs"] == body["mm_processor_kwargs"]
146
+ content = prepared["messages"][0]["content"]
147
+ assert [item["type"] for item in content] == ["audio_url", second_type]
148
+ assert content[0]["uuid"] == "user-audio"
149
+ for item, encoded in zip(content, seen):
150
+ assert item == encoded["messages"][0]["content"][0]
151
+ params = prepared["ec_transfer_params"]
152
+ for mm_hash, entry in reported.items():
153
+ assert params[mm_hash] == entry
154
+ if transfer == "mooncake":
155
+ assert params["ec_items"] == [
156
+ {
157
+ "mm_hash": mm_hash,
158
+ "transfer_id": encoded["ec_transfer_params"]["ec_items"][0][
159
+ "transfer_id"
160
+ ],
161
+ }
162
+ for mm_hash, encoded in zip(reported, seen)
163
+ ]
164
+
165
+
166
@pytest.mark.parametrize("no_rewrite", [False, True])
167
@pytest.mark.parametrize(
168
"batch_size, expected_sizes", [(0, [3, 3]), (1, [1] * 6), (2, [2, 2, 1, 1])]
@@ -509,7 +585,11 @@ def test_video_audio_fallback(proxy, monkeypatch, no_rewrite, transfer):
585
original = copy.deepcopy(body)
586
# Collector indices describe two processed features from one video item.
587
video_metadata = collect_ec_item_metadata(
512
- [SimpleNamespace(identifier=key, data=None) for key in ("video", "audio")], None
588
+ [
589
+ SimpleNamespace(identifier=key, modality=key, data=None)
590
+ for key in ("video", "audio")
591
+ ],
592
+ SimpleNamespace(fields_for=lambda _: set()),
593
)
594
if transfer == "handle":
595
video_metadata["audio"]["transfer_id"] = "reservation"
vllm/distributed/ec_transfer/ec_connector/utils.py
+8
-6
@@ -60,24 +60,26 @@ def collect_ec_item_metadata(
60
61
Keyed by mm_hash, each entry carries a `metadata` dict with whatever
62
placeholder fields `resolver` says this model needs published for its
63
- modality, so a consumer can skip the image transform. `item_indices`
63
+ modality, so a consumer can skip media preprocessing. `item_indices`
64
identifies every occurrence in `mm_features`, including repeated hashes.
65
- `data` is None for items served from the processor cache, so the metadata is
66
- unavailable here and the consumer has to fall back to processing the
67
- media itself. A connector that also has transfer coordinates to report
65
+ Audio token counts come from the placeholder; other metadata requires
66
+ `data` and falls back to raw media when it is unavailable on cache hits.
67
+ A connector that also has transfer coordinates to report
68
(e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside
69
`metadata`, not into it.
70
"""
71
items: dict[str, dict[str, Any]] = {}
72
for index, feature in enumerate(mm_features):
73
metadata: dict[str, Any] = {}
74
+ wanted = resolver.fields_for(feature.modality)
75
+ if "audio_num_tokens" in wanted:
76
+ metadata["audio_num_tokens"] = [feature.mm_position.get_num_embeds()]
77
if feature.data is not None:
75
- wanted = resolver.fields_for(feature.modality)
78
for key, value in feature.data.get_data().items():
79
if key not in wanted:
80
continue
81
if isinstance(value, torch.Tensor):
80
- metadata[key] = value.tolist()
82
+ metadata[key] = torch.atleast_1d(value).tolist()
83
elif is_list_of(value, (int, float)):
84
# Some metadata (e.g. Qwen3-VL video timestamps) is
85
# produced as a plain list rather than a tensor.
vllm/model_executor/models/audioflamingo3.py
+13
-2
@@ -197,6 +197,7 @@ class AudioFlamingo3ProcessingInfo(BaseProcessingInfo):
197
return AudioFlamingo3MultiModalDataParser(
198
target_sr=feature_extractor.sampling_rate,
199
expected_hidden_size=self._get_expected_hidden_size(),
200
+ allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
201
)
202
203
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
@@ -237,6 +238,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
238
if chunk_counts is not None:
239
return dict(
240
audio_embeds=MultiModalFieldConfig.batched("audio"),
241
+ audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
242
input_features=MultiModalFieldConfig.flat_from_sizes(
243
"audio", chunk_counts, dim=0
244
),
@@ -247,6 +249,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]):
249
)
250
return dict(
251
audio_embeds=MultiModalFieldConfig.batched("audio"),
252
+ audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
253
input_features=MultiModalFieldConfig.batched("audio"),
254
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
255
chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
@@ -354,11 +357,15 @@ def _count_audio_tokens_from_mask(
357
358
359
class AudioFlamingo3MultiModalDataParser(MultiModalDataParser):
360
+ embedding_fields = {
361
+ "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
362
+ }
363
+
364
def _parse_audio_data(
365
self,
366
data: dict[str, torch.Tensor] | ModalityData[Any],
367
) -> ModalityDataItems[Any, Any] | None:
361
- if isinstance(data, dict):
368
+ if isinstance(data, dict) and "audio_num_tokens" not in data:
369
return DictEmbeddingItems(
370
data,
371
modality="audio",
@@ -442,7 +449,11 @@ class AudioFlamingo3MultiModalProcessor(
449
chunk_counts = out_mm_data.get("chunk_counts")
450
451
def get_replacement_audioflamingo3(item_idx: int):
445
- if feature_attention_mask is not None:
452
+ if "audio_num_tokens" in out_mm_data:
453
+ counts = out_mm_data["audio_num_tokens"]
454
+ assert isinstance(counts, torch.Tensor)
455
+ num_features = int(counts[item_idx])
456
+ elif feature_attention_mask is not None:
457
num_features = _count_audio_tokens_from_mask(
458
feature_attention_mask,
459
chunk_counts,
vllm/model_executor/models/qwen2_5_omni_thinker.py
+15
-2
@@ -282,7 +282,7 @@ def create_qwen2_5_omni_thinker_field_factory(
282
def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]):
283
audio_feature_lengths = hf_inputs.get(
284
"audio_feature_lengths", torch.empty((0,))
285
- )
285
+ ).flatten()
286
287
image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3)))
288
image_pixel_grid_sizes = image_grid_thw.prod(-1)
@@ -332,6 +332,14 @@ def create_qwen2_5_omni_thinker_field_factory(
332
333
334
class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
335
+ embedding_fields = {
336
+ **Qwen2VLMultiModalDataParser.embedding_fields,
337
+ "audio": {
338
+ "input_audio_features": "values",
339
+ "audio_feature_lengths": "metadata",
340
+ },
341
+ }
342
+
343
def __init__(self, spatial_merge_size: int, *args, **kwargs):
344
self._spatial_merge_size = spatial_merge_size
345
super().__init__(self._spatial_merge_size, *args, **kwargs)
@@ -341,10 +349,15 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser):
349
data: dict[str, torch.Tensor] | ModalityData[ImageItem],
350
) -> ModalityDataItems[Any, Any]:
351
if isinstance(data, dict):
352
+ required, optional = self.embedding_field_sets("audio")
353
+ data = dict(data)
354
+ if "audio_feature_lengths" in data:
355
+ data["audio_feature_lengths"] = data["audio_feature_lengths"].flatten()
356
return DictEmbeddingItems(
357
data,
358
modality="audio",
347
- required_fields={"input_audio_features", "audio_feature_lengths"},
359
+ required_fields=required,
360
+ optional_fields=optional,
361
fields_factory=create_qwen2_5_omni_thinker_field_factory(
362
self._spatial_merge_size
363
),
vllm/model_executor/models/qwen2_audio.py
+10
-2
@@ -128,17 +128,22 @@ def _get_feat_extract_output_lengths(input_lengths: torch.Tensor):
128
def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]):
129
return dict(
130
audio_embeds=MultiModalFieldConfig.batched("audio"),
131
+ audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
132
input_features=MultiModalFieldConfig.batched("audio"),
133
feature_attention_mask=MultiModalFieldConfig.batched("audio"),
134
)
135
136
137
class Qwen2AudioMultiModalDataParser(MultiModalDataParser):
138
+ embedding_fields = {
139
+ "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
140
+ }
141
+
142
def _parse_audio_data(
143
self,
144
data: dict[str, torch.Tensor] | ModalityData[AudioItem],
145
) -> ModalityDataItems[Any, Any] | None:
141
- if isinstance(data, dict):
146
+ if isinstance(data, dict) and "audio_num_tokens" not in data:
147
return DictEmbeddingItems(
148
data,
149
modality="audio",
@@ -169,6 +174,7 @@ class Qwen2AudioProcessingInfo(BaseProcessingInfo):
174
target_sr=feature_extractor.sampling_rate,
175
target_channels=self.get_target_channels(),
176
expected_hidden_size=self._get_expected_hidden_size(),
177
+ allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
178
)
179
180
def get_target_channels(self) -> int:
@@ -279,7 +285,9 @@ class Qwen2AudioMultiModalProcessor(BaseMultiModalProcessor[Qwen2AudioProcessing
285
audio_output_lengths = audio_output_lens.tolist()
286
287
def get_replacement_qwen2_audio(item_idx: int):
282
- if audio_output_lengths:
288
+ if "audio_num_tokens" in out_mm_data:
289
+ num_features = int(out_mm_data["audio_num_tokens"][item_idx])
290
+ elif audio_output_lengths:
291
num_features = audio_output_lengths[item_idx]
292
else:
293
audio_embeds = out_mm_data["audio_embeds"][item_idx]
vllm/model_executor/models/qwen3_omni_moe_thinker.py
+17
@@ -100,6 +100,7 @@ from .qwen2_5_omni_thinker import (
100
Qwen2_5OmniAudioFeatureInputs,
101
Qwen2_5OmniConditionalGenerationMixin,
102
Qwen2_5OmniThinkerDummyInputsBuilder,
103
+ Qwen2_5OmniThinkerMultiModalDataParser,
104
Qwen2_5OmniThinkerMultiModalProcessor,
105
check_interleaved_audio_video,
106
merge_interleaved_embeddings,
@@ -1152,6 +1153,13 @@ class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM):
1153
)
1154
1155
1156
+class Qwen3OmniMoeThinkerMultiModalDataParser(Qwen2_5OmniThinkerMultiModalDataParser):
1157
+ # Keep the existing raw-media fallback for vision inputs.
1158
+ embedding_fields = {
1159
+ "audio": Qwen2_5OmniThinkerMultiModalDataParser.embedding_fields["audio"],
1160
+ }
1161
+
1162
+
1163
class Qwen3OmniMoeThinkerProcessingInfo(
1164
Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo
1165
):
@@ -1178,6 +1186,15 @@ class Qwen3OmniMoeThinkerProcessingInfo(
1186
assert isinstance(feature_extractor, WhisperFeatureExtractor)
1187
return feature_extractor
1188
1189
+ def get_data_parser(self):
1190
+ return Qwen3OmniMoeThinkerMultiModalDataParser(
1191
+ spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size,
1192
+ target_sr=self.get_feature_extractor().sampling_rate,
1193
+ target_channels=self.get_target_channels(),
1194
+ expected_hidden_size=self._get_expected_hidden_size(),
1195
+ allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
1196
+ )
1197
+
1198
def get_supported_mm_limits(self) -> Mapping[str, int | None]:
1199
return {"audio": None, "image": None, "video": None}
1200
vllm/model_executor/models/ultravox.py
+11
-1
@@ -114,6 +114,12 @@ UltravoxAudioInputs: TypeAlias = (
114
)
115
116
117
+class UltravoxMultiModalDataParser(MultiModalDataParser):
118
+ embedding_fields = {
119
+ "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"},
120
+ }
121
+
122
+
123
class UltravoxProcessingInfo(BaseProcessingInfo):
124
def get_hf_processor(self, **kwargs: object) -> ProcessorMixin:
125
config = self.ctx.model_config.hf_config
@@ -145,10 +151,11 @@ class UltravoxProcessingInfo(BaseProcessingInfo):
151
def get_data_parser(self):
152
feature_extractor = self.get_feature_extractor()
153
148
- return MultiModalDataParser(
154
+ return UltravoxMultiModalDataParser(
155
target_sr=feature_extractor.sampling_rate,
156
target_channels=self.get_target_channels(),
157
expected_hidden_size=self._get_expected_hidden_size(),
158
+ allow_missing_mm_embeddings=self.allow_missing_mm_embeddings,
159
)
160
161
def get_target_channels(self) -> int:
@@ -247,6 +254,7 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo
254
# num_chunks can convert audio_chunked to audio batch dimension
255
audio_num_chunks=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
256
audio_embeds=MultiModalFieldConfig.batched("audio"),
257
+ audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True),
258
)
259
260
def _get_prompt_updates(
@@ -273,6 +281,8 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo
281
)
282
283
def get_replacement_ultravox(item_idx: int):
284
+ if "audio_num_tokens" in out_mm_data:
285
+ return [replacement_id] * int(out_mm_data["audio_num_tokens"][item_idx])
286
start = chunks_start_idx[item_idx]
287
end = chunks_start_idx[item_idx + 1]
288
audio_token_len = out_mm_data["audio_token_len"][start:end].sum()
vllm/multimodal/parse.py
+43
-1
@@ -664,9 +664,46 @@ class MultiModalDataParser:
664
665
assert_never(video)
666
667
+ def _parse_audio_embedding_data(
668
+ self, data: dict[str, torch.Tensor]
669
+ ) -> DictEmbeddingItems:
670
+ counts = data.get("audio_num_tokens")
671
+ if (
672
+ not isinstance(counts, torch.Tensor)
673
+ or counts.dtype not in (torch.int32, torch.int64)
674
+ or counts.ndim not in (1, 2)
675
+ or (counts.ndim == 2 and counts.shape[1] != 1)
676
+ or not torch.all(counts > 0)
677
+ ):
678
+ raise ValueError(
679
+ "audio_num_tokens must contain one positive integer per audio"
680
+ )
681
+
682
+ counts = counts.flatten()
683
+ if "audio_embeds" in data:
684
+ embeds = data["audio_embeds"]
685
+ if len(embeds) != len(counts) or any(
686
+ len(embedding) != count for embedding, count in zip(embeds, counts)
687
+ ):
688
+ raise ValueError("audio_num_tokens does not match audio_embeds")
689
+
690
+ required, optional = self.embedding_field_sets("audio")
691
+ return DictEmbeddingItems(
692
+ {**data, "audio_num_tokens": counts},
693
+ modality="audio",
694
+ required_fields=required,
695
+ optional_fields=optional,
696
+ fields_factory=lambda _: {
697
+ "audio_embeds": MultiModalFieldConfig.batched("audio"),
698
+ "audio_num_tokens": MultiModalFieldConfig.batched(
699
+ "audio", keep_on_cpu=True
700
+ ),
701
+ },
702
+ )
703
+
704
def _parse_audio_data(
705
self,
669
- data: ModalityData[AudioItem],
706
+ data: dict[str, torch.Tensor] | ModalityData[AudioItem],
707
) -> ModalityDataItems[Any, Any] | None:
708
if data is None:
709
return None
@@ -674,6 +711,11 @@ class MultiModalDataParser:
711
if self.is_embeddings(data):
712
return AudioEmbeddingItems(data, self.expected_hidden_size)
713
714
+ if isinstance(data, dict) and "audio_num_tokens" in self.embedding_fields.get(
715
+ "audio", {}
716
+ ):
717
+ return self._parse_audio_embedding_data(data)
718
+
719
data_items: list[AudioItem | None]
720
if (
721
(is_list_of(data, float) and len(data) > 0)