[EPD] Support metadata-only audio inputs

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

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

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)

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 @@
     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,

tests/multimodal/test_parse.py

@@ -15,6 +15,45 @@
 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",
     [

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."""

tests/v1/ec_connector/unit/test_epd_proxy_retry.py

@@ -87,6 +87,82 @@ async def post(self, url, data=None, headers=None):
         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"