From 3ee2906f47afdfc1d792e04995c31eb6ecf2bdef Mon Sep 17 00:00:00 2001 From: Tianyu Guo Date: Wed, 23 Sep 2026 09:03:50 +0800 Subject: [PATCH] [EPD] Support metadata-only audio inputs (#57887) Signed-off-by: Tianyu Guo Co-authored-by: Cyrus Leung --- docs/features/disagg_encoder.md | 13 +++ .../disaggregated_encoder/disagg_epd_proxy.py | 45 +++++++--- tests/model_executor/test_qwen3_omni.py | 5 ++ .../multimodal/processing/test_common.py | 83 ++++++++++++++++++- tests/multimodal/test_parse.py | 39 +++++++++ .../unit/test_ec_example_connector.py | 16 ++++ .../ec_connector/unit/test_epd_proxy_retry.py | 82 +++++++++++++++++- .../ec_transfer/ec_connector/utils.py | 14 ++-- vllm/model_executor/models/audioflamingo3.py | 15 +++- .../models/qwen2_5_omni_thinker.py | 17 +++- vllm/model_executor/models/qwen2_audio.py | 12 ++- .../models/qwen3_omni_moe_thinker.py | 17 ++++ vllm/model_executor/models/ultravox.py | 12 ++- vllm/multimodal/parse.py | 44 +++++++++- 14 files changed, 387 insertions(+), 27 deletions(-) diff --git a/docs/features/disagg_encoder.md b/docs/features/disagg_encoder.md index 383e7d45ff..ac304345fd 100644 --- a/docs/features/disagg_encoder.md +++ b/docs/features/disagg_encoder.md @@ -61,6 +61,19 @@ and compares the disaggregated outputs against it. Set `MODEL` to use another mo or `MOONCAKE_EC_PROTOCOL=rdma` to use RDMA on supported hardware. See the script for additional configuration options. +### Audio inputs + +The Python EPD proxy also rewrites audio in `/v1/chat/completions` into +metadata-only references for Qwen2-Audio, AudioFlamingo3, Ultravox, +Qwen2.5-Omni, and Qwen3-Omni. The encoder publishes `audio_num_tokens` +(the placeholder token count of each audio) together with any feature +lengths the model needs; the consumer uses them to reconstruct the audio +placeholders and receives the embeddings through its EC connector without +repeating audio preprocessing. + +This covers pure audio inputs, not video with an embedded audio track. It does +not add `/v1/audio/transcriptions` or realtime routes to the example proxy. + --- ## 3 Test Script diff --git a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py index abae7eb65d..4715a15f86 100644 --- a/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py +++ b/examples/disaggregated/disaggregated_encoder/disagg_epd_proxy.py @@ -480,6 +480,17 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict: the grid its processor actually produced), so the grid is never re-derived here -- a second derivation could disagree with the encoder's. """ + mm_items = extract_mm_items(req_data) + # Audio processors may consume placeholders from other modalities too. + # Keep the whole request raw if any item cannot be rewritten safely. + control_fields = {"mm_hash", "ec_mm_hash", "transfer_id"} + raw_audio_request = any( + item["type"] in {"audio_url", "input_audio"} for item in mm_items + ) and any( + not item_meta.get(i, {}).get("mm_hash") + or not (item_meta.get(i, {}).keys() - control_fields) + for i in range(len(mm_items)) + ) rewritten = 0 transfer_items = [] idx = 0 @@ -499,6 +510,18 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict: item_uuid = meta.pop("mm_hash", None) ec_mm_hash = meta.pop("ec_mm_hash", None) or item_uuid transfer_id = meta.pop("transfer_id", None) + if raw_audio_request: + new_content.append( + { + **item, + "uuid": item_uuid or item.get("uuid") or content_uuid(item), + } + ) + if transfer_id is not None: + transfer_items.append( + {"mm_hash": ec_mm_hash, "transfer_id": transfer_id} + ) + continue # Whatever keys the encoder reported are the metadata its model # declared as needed to size the placeholder range; the proxy does # not need to know their names. @@ -532,9 +555,10 @@ def rewrite_for_decode(req_data: dict, item_meta: dict[int, dict]) -> dict: rewritten += 1 new_messages.append({**msg, "content": new_content}) - if not rewritten: + if not rewritten and not raw_audio_request: return req_data - logger.info("Rewrote %d media item(s) as metadata references", rewritten) + if rewritten: + logger.info("Rewrote %d media item(s) as metadata references", rewritten) rewritten_request = {**req_data, "messages": new_messages} if transfer_items: ec_transfer_params = dict(req_data.get("ec_transfer_params") or {}) @@ -632,7 +656,7 @@ async def fanout_encoder_primer( content = [] for idx in indices: item = mm_items[idx] - item_uuid = None if NO_REWRITE else content_uuid(item) + item_uuid = None if NO_REWRITE else (item.get("uuid") or content_uuid(item)) if item_uuid is not None: item_uuids[idx] = item_uuid item_transfer_ids[idx] = uuid.uuid4().hex @@ -756,14 +780,13 @@ async def fanout_encoder_primer( raise HTTPException(502, "Encoder metadata cannot be matched") ec_mm_hash, reported = matched metadata = reported.get("metadata") or {} - if metadata: - item_meta[idx] = { - **metadata, - "mm_hash": item_uuids.get(idx, ec_mm_hash), - "ec_mm_hash": ec_mm_hash, - } - if idx in item_transfer_ids: - item_meta[idx]["transfer_id"] = item_transfer_ids[idx] + item_meta[idx] = { + **metadata, + "mm_hash": item_uuids.get(idx, ec_mm_hash), + "ec_mm_hash": ec_mm_hash, + } + if idx in item_transfer_ids: + item_meta[idx]["transfer_id"] = item_transfer_ids[idx] # Whatever the encoder reported alongside `metadata` is the # connector's own handle on the published embedding (for NIXL, # peer_host/peer_port/size_bytes). The decoder's connector diff --git a/tests/model_executor/test_qwen3_omni.py b/tests/model_executor/test_qwen3_omni.py index 71c74863e9..cdf169601e 100644 --- a/tests/model_executor/test_qwen3_omni.py +++ b/tests/model_executor/test_qwen3_omni.py @@ -144,6 +144,11 @@ def test_qwen3_omni_get_updates_use_audio_in_video( # Create a mock context mock_ctx = Mock(spec=InputProcessingContext) mock_ctx.tokenizer = mock_tokenizer + # `model_config` is an instance attribute, so it is not covered by the + # spec; the data parser reads it for `allow_missing_mm_embeddings`. + mock_ctx.model_config = SimpleNamespace( + multimodal_config=SimpleNamespace(allow_missing_mm_embeddings=False) + ) # Create processing info info = Qwen3OmniMoeThinkerProcessingInfo(mock_ctx) diff --git a/tests/models/multimodal/processing/test_common.py b/tests/models/multimodal/processing/test_common.py index fa23956017..20a5121e60 100644 --- a/tests/models/multimodal/processing/test_common.py +++ b/tests/models/multimodal/processing/test_common.py @@ -1,13 +1,16 @@ # SPDX-License-Identifier: Apache-2.0 # SPDX-FileCopyrightText: Copyright contributors to the vLLM project +import json from collections.abc import Set as AbstractSet from functools import partial import numpy as np import pytest +import torch from PIL import Image +from tests.models.utils import build_model_context from vllm.config import ModelConfig from vllm.config.multimodal import ( AudioDummyOptions, @@ -16,10 +19,15 @@ from vllm.config.multimodal import ( MultiModalDummyOptions, VideoDummyOptions, ) +from vllm.distributed.ec_transfer.ec_connector.utils import ( + PlaceholderMetadataResolver, + collect_ec_item_metadata, +) +from vllm.entrypoints.chat_utils import _get_embeds_data from vllm.inputs import MultiModalDataDict, MultiModalInput from vllm.multimodal import MULTIMODAL_REGISTRY from vllm.multimodal.cache import MultiModalProcessorOnlyCache -from vllm.multimodal.inputs import batched_tensors_equal +from vllm.multimodal.inputs import MultiModalFeatureSpec, batched_tensors_equal from vllm.multimodal.processing import BaseMultiModalProcessor, InputProcessingContext from vllm.platforms import current_platform from vllm.tokenizers import TokenizerLike, cached_tokenizer_from_config @@ -193,6 +201,79 @@ def get_token_prompt( return inputs.prompt +@pytest.mark.parametrize( + "model_id,durations", + [ + ("Qwen/Qwen2-Audio-7B-Instruct", [1.0, 2.3]), + ("nvidia/audio-flamingo-3-hf", [31.0]), + ("fixie-ai/ultravox-v0_5-llama-3_2-1b", [1.0, 31.0]), + ("Qwen/Qwen2.5-Omni-3B", [1.0, 2.3]), + ("Qwen/Qwen3-Omni-30B-A3B-Instruct", [1.0, 2.3]), + ], +) +def test_audio_metadata_only_roundtrip(model_id, durations, monkeypatch): + """EC metadata preserves the raw-audio prompt without running the HF processor.""" + ctx = build_model_context(model_id, limit_mm_per_prompt={"audio": len(durations)}) + mm_config = ctx.model_config.multimodal_config + mm_config.enable_mm_embeds = True + mm_config.allow_missing_mm_embeddings = True + processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + sample_rate = processor.info.get_feature_extractor().sampling_rate + audios = [ + np.zeros(int(duration * sample_rate), dtype=np.float32) + for duration in durations + ] + prompt = processor.dummy_inputs.get_dummy_text({"audio": len(audios)}) + raw = processor( + prompt, + mm_items=processor.info.parse_mm_data({"audio": audios}), + hf_processor_mm_kwargs={}, + ) + features = [ + MultiModalFeatureSpec( + data=item, + modality="audio", + identifier=str(index), + mm_position=position, + ) + for index, (item, position) in enumerate( + zip(raw["mm_kwargs"]["audio"], raw["mm_placeholders"]["audio"]) + ) + ] + published = collect_ec_item_metadata( + features, PlaceholderMetadataResolver(ctx.model_config) + ) + metadata = json.loads( + json.dumps([entry["metadata"] for entry in published.values()]) + ) + assert len(metadata) == len(audios) + assert all(metadata) + audio_data = _get_embeds_data("audio", metadata, processor) + + def unexpected_preprocessing(*args, **kwargs): + pytest.fail("The EC consumer must not preprocess audio") + + monkeypatch.setattr(processor, "_call_hf_processor", unexpected_preprocessing) + remote = processor( + prompt, + mm_items=processor.info.parse_mm_data({"audio": audio_data}), + hf_processor_mm_kwargs={}, + ) + assert remote["prompt_token_ids"] == raw["prompt_token_ids"] + assert remote["mm_placeholders"] == raw["mm_placeholders"] + # Omni also consumes these lengths when constructing M-RoPE positions. + if "audio_feature_lengths" in audio_data: + assert torch.equal( + remote["mm_kwargs"].get_data()["audio_feature_lengths"], + raw["mm_kwargs"].get_data()["audio_feature_lengths"], + ) + + mm_config.allow_missing_mm_embeddings = False + ordinary_processor = MULTIMODAL_REGISTRY.create_processor(ctx.model_config) + with pytest.raises(ValueError, match="should contain the fields"): + ordinary_processor.info.parse_mm_data({"audio": audio_data}) + + def random_vision_chunk( rng: np.random.RandomState, min_wh: int, diff --git a/tests/multimodal/test_parse.py b/tests/multimodal/test_parse.py index 2649f80777..c20677fc25 100644 --- a/tests/multimodal/test_parse.py +++ b/tests/multimodal/test_parse.py @@ -15,6 +15,45 @@ from vllm.multimodal.parse import ( H, W = 480, 640 +class AudioMetadataParser(MultiModalDataParser): + embedding_fields = { + "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"}, + } + + +@pytest.mark.parametrize("allow_missing", [False, True]) +def test_audio_metadata_requires_ec_consumer(allow_missing): + parser = AudioMetadataParser(allow_missing_mm_embeddings=allow_missing) + data = {"audio_num_tokens": torch.tensor([[3], [5]])} + if not allow_missing: + with pytest.raises(ValueError, match="audio_embeds"): + parser.parse_mm_data({"audio": data}) + else: + items = parser.parse_mm_data({"audio": data})["audio"] + assert len(items) == 2 + assert items.get(1)["audio_num_tokens"].item() == 5 + assert items.get_processor_data() == {} + + +@pytest.mark.parametrize("counts", [[0], [-1], [1.5], [True], [[1, 2]]]) +def test_audio_metadata_rejects_invalid_token_counts(counts): + parser = AudioMetadataParser(allow_missing_mm_embeddings=True) + with pytest.raises(ValueError, match="positive integer"): + parser.parse_mm_data({"audio": {"audio_num_tokens": torch.tensor(counts)}}) + + +def test_audio_metadata_checks_supplied_embedding_lengths(): + parser = AudioMetadataParser() + data = { + "audio_num_tokens": torch.tensor([3, 5]), + "audio_embeds": [torch.zeros(3, 8), torch.zeros(5, 8)], + } + assert len(parser.parse_mm_data({"audio": data})["audio"]) == 2 + data["audio_num_tokens"] = torch.tensor([3, 4]) + with pytest.raises(ValueError, match="does not match"): + parser.parse_mm_data({"audio": data}) + + @pytest.mark.parametrize( "image", [ diff --git a/tests/v1/ec_connector/unit/test_ec_example_connector.py b/tests/v1/ec_connector/unit/test_ec_example_connector.py index a6645fccf4..725bb99eb0 100644 --- a/tests/v1/ec_connector/unit/test_ec_example_connector.py +++ b/tests/v1/ec_connector/unit/test_ec_example_connector.py @@ -56,6 +56,22 @@ def test_metadata_preserves_repeated_item_positions(): } +def test_audio_metadata_survives_processor_cache_hit(): + resolver = Mock() + resolver.fields_for.return_value = {"audio_num_tokens"} + feature = MultiModalFeatureSpec( + data=None, + modality="audio", + identifier="audio", + mm_position=PlaceholderRange( + offset=0, length=4, is_embed=torch.tensor([False, True, True, False]) + ), + ) + assert collect_ec_item_metadata([feature], resolver) == { + "audio": {"metadata": {"audio_num_tokens": [2]}, "item_indices": [0]} + } + + @pytest.fixture def temp_storage(tmp_path): """Fixture providing temporary storage path.""" diff --git a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py index 8db22fe1b9..08bca0e40f 100644 --- a/tests/v1/ec_connector/unit/test_epd_proxy_retry.py +++ b/tests/v1/ec_connector/unit/test_epd_proxy_retry.py @@ -87,6 +87,82 @@ class _EncoderSession: return _EncoderResponse(self._replies.pop(0)) +@pytest.mark.parametrize("second_type", ["image_url", "audio_url", "input_audio"]) +@pytest.mark.parametrize("transfer", ["example", "nixl", "mooncake"]) +def test_partial_audio_metadata_falls_back_with_all_transfers( + proxy, monkeypatch, second_type, transfer +): + """Never mix metadata-only audio with raw media, or lose EC identities.""" + import copy + + seen = [] + reported = {} + + class Session: + async def post(self, url, data=None, headers=None): + body = msgspec.json.decode(data) + seen.append(body) + item = body["messages"][0]["content"][0] + mm_hash = item["uuid"] + "-processed" + entry: dict[str, Any] = { + "metadata": {"audio_feature_lengths": [100]} if len(seen) == 1 else {}, + "item_indices": [0], + } + if transfer == "nixl": + entry.update(peer_host="encoder", peer_port=1234, size_bytes=4096) + reported[mm_hash] = entry + return _EncoderResponse({mm_hash: entry}) + + second = ( + {"type": second_type, "input_audio": {"data": "large", "format": "wav"}} + if second_type == "input_audio" + else {"type": second_type, second_type: {"url": "large"}} + ) + body = { + "messages": [ + { + "role": "user", + "content": [ + { + "type": "audio_url", + "audio_url": {"url": "small"}, + "uuid": "user-audio", + }, + second, + ], + } + ], + "mm_processor_kwargs": {"sampling_rate": 16000}, + } + original = copy.deepcopy(body) + monkeypatch.setattr(proxy, "encode_session", Session()) + monkeypatch.setattr(proxy, "NO_REWRITE", False) + consumer = "tcp://consumer:1" if transfer == "mooncake" else None + prepared, _, _ = asyncio.run( + proxy.prepare_for_decode(body, "r", ["http://encoder"], "", consumer) + ) + assert body == original + assert prepared["mm_processor_kwargs"] == body["mm_processor_kwargs"] + content = prepared["messages"][0]["content"] + assert [item["type"] for item in content] == ["audio_url", second_type] + assert content[0]["uuid"] == "user-audio" + for item, encoded in zip(content, seen): + assert item == encoded["messages"][0]["content"][0] + params = prepared["ec_transfer_params"] + for mm_hash, entry in reported.items(): + assert params[mm_hash] == entry + if transfer == "mooncake": + assert params["ec_items"] == [ + { + "mm_hash": mm_hash, + "transfer_id": encoded["ec_transfer_params"]["ec_items"][0][ + "transfer_id" + ], + } + for mm_hash, encoded in zip(reported, seen) + ] + + @pytest.mark.parametrize("no_rewrite", [False, True]) @pytest.mark.parametrize( "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): original = copy.deepcopy(body) # Collector indices describe two processed features from one video item. video_metadata = collect_ec_item_metadata( - [SimpleNamespace(identifier=key, data=None) for key in ("video", "audio")], None + [ + SimpleNamespace(identifier=key, modality=key, data=None) + for key in ("video", "audio") + ], + SimpleNamespace(fields_for=lambda _: set()), ) if transfer == "handle": video_metadata["audio"]["transfer_id"] = "reservation" diff --git a/vllm/distributed/ec_transfer/ec_connector/utils.py b/vllm/distributed/ec_transfer/ec_connector/utils.py index 35d26bb44b..081d05068e 100644 --- a/vllm/distributed/ec_transfer/ec_connector/utils.py +++ b/vllm/distributed/ec_transfer/ec_connector/utils.py @@ -60,24 +60,26 @@ def collect_ec_item_metadata( Keyed by mm_hash, each entry carries a `metadata` dict with whatever placeholder fields `resolver` says this model needs published for its - modality, so a consumer can skip the image transform. `item_indices` + modality, so a consumer can skip media preprocessing. `item_indices` identifies every occurrence in `mm_features`, including repeated hashes. - `data` is None for items served from the processor cache, so the metadata is - unavailable here and the consumer has to fall back to processing the - media itself. A connector that also has transfer coordinates to report + Audio token counts come from the placeholder; other metadata requires + `data` and falls back to raw media when it is unavailable on cache hits. + A connector that also has transfer coordinates to report (e.g. NIXL peer_host/peer_port/size_bytes) merges those in alongside `metadata`, not into it. """ items: dict[str, dict[str, Any]] = {} for index, feature in enumerate(mm_features): metadata: dict[str, Any] = {} + wanted = resolver.fields_for(feature.modality) + if "audio_num_tokens" in wanted: + metadata["audio_num_tokens"] = [feature.mm_position.get_num_embeds()] if feature.data is not None: - wanted = resolver.fields_for(feature.modality) for key, value in feature.data.get_data().items(): if key not in wanted: continue if isinstance(value, torch.Tensor): - metadata[key] = value.tolist() + metadata[key] = torch.atleast_1d(value).tolist() elif is_list_of(value, (int, float)): # Some metadata (e.g. Qwen3-VL video timestamps) is # produced as a plain list rather than a tensor. diff --git a/vllm/model_executor/models/audioflamingo3.py b/vllm/model_executor/models/audioflamingo3.py index 4be996d278..edbe92101d 100644 --- a/vllm/model_executor/models/audioflamingo3.py +++ b/vllm/model_executor/models/audioflamingo3.py @@ -197,6 +197,7 @@ class AudioFlamingo3ProcessingInfo(BaseProcessingInfo): return AudioFlamingo3MultiModalDataParser( target_sr=feature_extractor.sampling_rate, expected_hidden_size=self._get_expected_hidden_size(), + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_supported_mm_limits(self) -> Mapping[str, int | None]: @@ -237,6 +238,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]): if chunk_counts is not None: return dict( audio_embeds=MultiModalFieldConfig.batched("audio"), + audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), input_features=MultiModalFieldConfig.flat_from_sizes( "audio", chunk_counts, dim=0 ), @@ -247,6 +249,7 @@ def _audioflamingo3_field_config(hf_inputs: Mapping[str, torch.Tensor]): ) return dict( audio_embeds=MultiModalFieldConfig.batched("audio"), + audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), input_features=MultiModalFieldConfig.batched("audio"), feature_attention_mask=MultiModalFieldConfig.batched("audio"), chunk_counts=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), @@ -354,11 +357,15 @@ def _count_audio_tokens_from_mask( class AudioFlamingo3MultiModalDataParser(MultiModalDataParser): + embedding_fields = { + "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"}, + } + def _parse_audio_data( self, data: dict[str, torch.Tensor] | ModalityData[Any], ) -> ModalityDataItems[Any, Any] | None: - if isinstance(data, dict): + if isinstance(data, dict) and "audio_num_tokens" not in data: return DictEmbeddingItems( data, modality="audio", @@ -442,7 +449,11 @@ class AudioFlamingo3MultiModalProcessor( chunk_counts = out_mm_data.get("chunk_counts") def get_replacement_audioflamingo3(item_idx: int): - if feature_attention_mask is not None: + if "audio_num_tokens" in out_mm_data: + counts = out_mm_data["audio_num_tokens"] + assert isinstance(counts, torch.Tensor) + num_features = int(counts[item_idx]) + elif feature_attention_mask is not None: num_features = _count_audio_tokens_from_mask( feature_attention_mask, chunk_counts, diff --git a/vllm/model_executor/models/qwen2_5_omni_thinker.py b/vllm/model_executor/models/qwen2_5_omni_thinker.py index ce7cb058b7..0613d9f8c5 100644 --- a/vllm/model_executor/models/qwen2_5_omni_thinker.py +++ b/vllm/model_executor/models/qwen2_5_omni_thinker.py @@ -282,7 +282,7 @@ def create_qwen2_5_omni_thinker_field_factory( def _qwen2_5_omni_thinker_field_config(hf_inputs: Mapping[str, torch.Tensor]): audio_feature_lengths = hf_inputs.get( "audio_feature_lengths", torch.empty((0,)) - ) + ).flatten() image_grid_thw = hf_inputs.get("image_grid_thw", torch.empty((0, 3))) image_pixel_grid_sizes = image_grid_thw.prod(-1) @@ -332,6 +332,14 @@ def create_qwen2_5_omni_thinker_field_factory( class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser): + embedding_fields = { + **Qwen2VLMultiModalDataParser.embedding_fields, + "audio": { + "input_audio_features": "values", + "audio_feature_lengths": "metadata", + }, + } + def __init__(self, spatial_merge_size: int, *args, **kwargs): self._spatial_merge_size = spatial_merge_size super().__init__(self._spatial_merge_size, *args, **kwargs) @@ -341,10 +349,15 @@ class Qwen2_5OmniThinkerMultiModalDataParser(Qwen2VLMultiModalDataParser): data: dict[str, torch.Tensor] | ModalityData[ImageItem], ) -> ModalityDataItems[Any, Any]: if isinstance(data, dict): + required, optional = self.embedding_field_sets("audio") + data = dict(data) + if "audio_feature_lengths" in data: + data["audio_feature_lengths"] = data["audio_feature_lengths"].flatten() return DictEmbeddingItems( data, modality="audio", - required_fields={"input_audio_features", "audio_feature_lengths"}, + required_fields=required, + optional_fields=optional, fields_factory=create_qwen2_5_omni_thinker_field_factory( self._spatial_merge_size ), diff --git a/vllm/model_executor/models/qwen2_audio.py b/vllm/model_executor/models/qwen2_audio.py index b643fe4227..dca856c45f 100644 --- a/vllm/model_executor/models/qwen2_audio.py +++ b/vllm/model_executor/models/qwen2_audio.py @@ -128,17 +128,22 @@ def _get_feat_extract_output_lengths(input_lengths: torch.Tensor): def _qwen2audio_field_config(hf_inputs: Mapping[str, torch.Tensor]): return dict( audio_embeds=MultiModalFieldConfig.batched("audio"), + audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), input_features=MultiModalFieldConfig.batched("audio"), feature_attention_mask=MultiModalFieldConfig.batched("audio"), ) class Qwen2AudioMultiModalDataParser(MultiModalDataParser): + embedding_fields = { + "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"}, + } + def _parse_audio_data( self, data: dict[str, torch.Tensor] | ModalityData[AudioItem], ) -> ModalityDataItems[Any, Any] | None: - if isinstance(data, dict): + if isinstance(data, dict) and "audio_num_tokens" not in data: return DictEmbeddingItems( data, modality="audio", @@ -169,6 +174,7 @@ class Qwen2AudioProcessingInfo(BaseProcessingInfo): target_sr=feature_extractor.sampling_rate, target_channels=self.get_target_channels(), expected_hidden_size=self._get_expected_hidden_size(), + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_target_channels(self) -> int: @@ -279,7 +285,9 @@ class Qwen2AudioMultiModalProcessor(BaseMultiModalProcessor[Qwen2AudioProcessing audio_output_lengths = audio_output_lens.tolist() def get_replacement_qwen2_audio(item_idx: int): - if audio_output_lengths: + if "audio_num_tokens" in out_mm_data: + num_features = int(out_mm_data["audio_num_tokens"][item_idx]) + elif audio_output_lengths: num_features = audio_output_lengths[item_idx] else: audio_embeds = out_mm_data["audio_embeds"][item_idx] diff --git a/vllm/model_executor/models/qwen3_omni_moe_thinker.py b/vllm/model_executor/models/qwen3_omni_moe_thinker.py index 85d0bcaad0..db3e8936ec 100755 --- a/vllm/model_executor/models/qwen3_omni_moe_thinker.py +++ b/vllm/model_executor/models/qwen3_omni_moe_thinker.py @@ -100,6 +100,7 @@ from .qwen2_5_omni_thinker import ( Qwen2_5OmniAudioFeatureInputs, Qwen2_5OmniConditionalGenerationMixin, Qwen2_5OmniThinkerDummyInputsBuilder, + Qwen2_5OmniThinkerMultiModalDataParser, Qwen2_5OmniThinkerMultiModalProcessor, check_interleaved_audio_video, merge_interleaved_embeddings, @@ -1152,6 +1153,13 @@ class Qwen3MoeLLMForCausalLM(Qwen3MoeForCausalLM): ) +class Qwen3OmniMoeThinkerMultiModalDataParser(Qwen2_5OmniThinkerMultiModalDataParser): + # Keep the existing raw-media fallback for vision inputs. + embedding_fields = { + "audio": Qwen2_5OmniThinkerMultiModalDataParser.embedding_fields["audio"], + } + + class Qwen3OmniMoeThinkerProcessingInfo( Qwen2AudioProcessingInfo, Qwen2_5_VLProcessingInfo ): @@ -1178,6 +1186,15 @@ class Qwen3OmniMoeThinkerProcessingInfo( assert isinstance(feature_extractor, WhisperFeatureExtractor) return feature_extractor + def get_data_parser(self): + return Qwen3OmniMoeThinkerMultiModalDataParser( + spatial_merge_size=self.get_hf_config().vision_config.spatial_merge_size, + target_sr=self.get_feature_extractor().sampling_rate, + target_channels=self.get_target_channels(), + expected_hidden_size=self._get_expected_hidden_size(), + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, + ) + def get_supported_mm_limits(self) -> Mapping[str, int | None]: return {"audio": None, "image": None, "video": None} diff --git a/vllm/model_executor/models/ultravox.py b/vllm/model_executor/models/ultravox.py index bb92833f77..a454ba5962 100644 --- a/vllm/model_executor/models/ultravox.py +++ b/vllm/model_executor/models/ultravox.py @@ -114,6 +114,12 @@ UltravoxAudioInputs: TypeAlias = ( ) +class UltravoxMultiModalDataParser(MultiModalDataParser): + embedding_fields = { + "audio": {"audio_embeds": "values", "audio_num_tokens": "metadata"}, + } + + class UltravoxProcessingInfo(BaseProcessingInfo): def get_hf_processor(self, **kwargs: object) -> ProcessorMixin: config = self.ctx.model_config.hf_config @@ -145,10 +151,11 @@ class UltravoxProcessingInfo(BaseProcessingInfo): def get_data_parser(self): feature_extractor = self.get_feature_extractor() - return MultiModalDataParser( + return UltravoxMultiModalDataParser( target_sr=feature_extractor.sampling_rate, target_channels=self.get_target_channels(), expected_hidden_size=self._get_expected_hidden_size(), + allow_missing_mm_embeddings=self.allow_missing_mm_embeddings, ) def get_target_channels(self) -> int: @@ -247,6 +254,7 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo # num_chunks can convert audio_chunked to audio batch dimension audio_num_chunks=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), audio_embeds=MultiModalFieldConfig.batched("audio"), + audio_num_tokens=MultiModalFieldConfig.batched("audio", keep_on_cpu=True), ) def _get_prompt_updates( @@ -273,6 +281,8 @@ class UltravoxMultiModalProcessor(BaseMultiModalProcessor[UltravoxProcessingInfo ) def get_replacement_ultravox(item_idx: int): + if "audio_num_tokens" in out_mm_data: + return [replacement_id] * int(out_mm_data["audio_num_tokens"][item_idx]) start = chunks_start_idx[item_idx] end = chunks_start_idx[item_idx + 1] audio_token_len = out_mm_data["audio_token_len"][start:end].sum() diff --git a/vllm/multimodal/parse.py b/vllm/multimodal/parse.py index 8d7b5aa9b6..98f8a8b2f3 100644 --- a/vllm/multimodal/parse.py +++ b/vllm/multimodal/parse.py @@ -664,9 +664,46 @@ class MultiModalDataParser: assert_never(video) + def _parse_audio_embedding_data( + self, data: dict[str, torch.Tensor] + ) -> DictEmbeddingItems: + counts = data.get("audio_num_tokens") + if ( + not isinstance(counts, torch.Tensor) + or counts.dtype not in (torch.int32, torch.int64) + or counts.ndim not in (1, 2) + or (counts.ndim == 2 and counts.shape[1] != 1) + or not torch.all(counts > 0) + ): + raise ValueError( + "audio_num_tokens must contain one positive integer per audio" + ) + + counts = counts.flatten() + if "audio_embeds" in data: + embeds = data["audio_embeds"] + if len(embeds) != len(counts) or any( + len(embedding) != count for embedding, count in zip(embeds, counts) + ): + raise ValueError("audio_num_tokens does not match audio_embeds") + + required, optional = self.embedding_field_sets("audio") + return DictEmbeddingItems( + {**data, "audio_num_tokens": counts}, + modality="audio", + required_fields=required, + optional_fields=optional, + fields_factory=lambda _: { + "audio_embeds": MultiModalFieldConfig.batched("audio"), + "audio_num_tokens": MultiModalFieldConfig.batched( + "audio", keep_on_cpu=True + ), + }, + ) + def _parse_audio_data( self, - data: ModalityData[AudioItem], + data: dict[str, torch.Tensor] | ModalityData[AudioItem], ) -> ModalityDataItems[Any, Any] | None: if data is None: return None @@ -674,6 +711,11 @@ class MultiModalDataParser: if self.is_embeddings(data): return AudioEmbeddingItems(data, self.expected_hidden_size) + if isinstance(data, dict) and "audio_num_tokens" in self.embedding_fields.get( + "audio", {} + ): + return self._parse_audio_embedding_data(data) + data_items: list[AudioItem | None] if ( (is_list_of(data, float) and len(data) > 0)