@tduncan / Nvidia-Optimization-Stack-VLLM Mirror · read-only / commits / 0730414999

Fork 0

A high-throughput and memory-efficient inference and serving engine for LLMs

Mirrored from https://github.com/tduncan-lum/Nvidia-Optimization-Stack---VLLM.git

Synced Nov 12, 19:45 UTC

[Core] Add audio_embeds support to chat completions (#29059)

Signed-off-by: Jeremy Teboul <jeremyteboul@fb.com> Co-authored-by: Jeremy Teboul <jeremyteboul@fb.com>

jeremyteboul committed Nov 20, 2025 at 19:39 UTC 0730414999343e722590ace615d5814c7e5b6827

5 files changed +360 -3

docs/features/multimodal_inputs.md

+32

@@ -365,6 +365,8 @@ You must enable this feature via `enable_mm_embeds=True`.

365 The vLLM engine may crash if incorrect shape of embeddings is passed.

366 Only enable this flag for trusted users!

367

368 +#### Image Embeddings

369 +

370 ??? code

371

372 ```python

@@ -441,6 +443,36 @@ For Qwen2-VL and MiniCPM-V, we accept additional parameters alongside the embedd

443 print(generated_text)

444 ```

445

446 +#### Audio Embeddings

447 +

448 +You can pass pre-computed audio embeddings similar to image embeddings:

449 +

450 +??? code

451 +

452 + ```python

453 + from vllm import LLM

454 + import torch

455 +

456 + # Enable audio embeddings support

457 + llm = LLM(model="fixie-ai/ultravox-v0_5-llama-3_2-1b", enable_mm_embeds=True)

458 +

459 + # Refer to the HuggingFace repo for the correct format to use

460 + prompt = "USER: <audio>\nWhat is in this audio?\nASSISTANT:"

461 +

462 + # Load pre-computed audio embeddings

463 + # torch.Tensor of shape (1, audio_feature_size, hidden_size of LM)

464 + audio_embeds = torch.load(...)

465 +

466 + outputs = llm.generate({

467 + "prompt": prompt,

468 + "multi_modal_data": {"audio": audio_embeds},

469 + })

470 +

471 + for o in outputs:

472 + generated_text = o.outputs[0].text

473 + print(generated_text)

474 + ```

475 +

476 ## Online Serving

477

478 Our OpenAI-compatible server accepts multi-modal data via the [Chat Completions API](https://platform.openai.com/docs/api-reference/chat). Media inputs also support optional UUIDs users can provide to uniquely identify each media, which is used to cache the media results across requests.

tests/entrypoints/test_chat_utils.py

+145

@@ -103,6 +103,19 @@ def qwen2_audio_model_config():

103 )

104

105

106 +@pytest.fixture(scope="function")

107 +def audio_embeds_model_config():

108 + return ModelConfig(

109 + QWEN2AUDIO_MODEL_ID,

110 + runner="generate",

111 + trust_remote_code=True,

112 + limit_mm_per_prompt={

113 + "audio": 2,

114 + },

115 + enable_mm_embeds=True,

116 + )

117 +

118 +

119 @pytest.fixture(scope="module")

120 def qwen2_audio_tokenizer():

121 return get_tokenizer(QWEN2AUDIO_MODEL_ID)

@@ -843,6 +856,138 @@ def test_parse_chat_messages_empty_image_embeds_with_uuid(

856 _assert_mm_uuids(mm_uuids, 1, expected_uuids=[uuid])

857

858

859 +def test_parse_chat_messages_empty_audio_embeds_with_uuid(

860 + audio_embeds_model_config,

861 + qwen2_audio_tokenizer,

862 +):

863 + """Test audio_embeds with UUID (no actual embeds data)."""

864 + uuid = "test-audio-uuid-123"

865 +

866 + conversation, mm_data, mm_uuids = parse_chat_messages(

867 + [

868 + {

869 + "role": "user",

870 + "content": [

871 + {"type": "text", "text": "Describe this audio"},

872 + {"type": "audio_embeds", "audio_embeds": None, "uuid": uuid},

873 + ],

874 + }

875 + ],

876 + audio_embeds_model_config,

877 + qwen2_audio_tokenizer,

878 + content_format="string",

879 + )

880 +

881 + # Should have audio in mm_data as None (UUID provided)

882 + assert mm_data is not None

883 + assert "audio" in mm_data

884 + assert mm_data["audio"] is None

885 + # UUID should be recorded

886 + assert mm_uuids is not None

887 + assert "audio" in mm_uuids

888 + _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[uuid])

889 +

890 +

891 +def test_parse_chat_messages_audio_embeds_with_string(

892 + audio_embeds_model_config,

893 + qwen2_audio_tokenizer,

894 +):

895 + """Test audio_embeds with base64 string embedding data."""

896 + import base64

897 + import io

898 +

899 + import torch

900 +

901 + # Create a sample audio embedding tensor

902 + audio_embedding = torch.randn(1, 128, 768)

903 +

904 + # Encode it as base64

905 + buffer = io.BytesIO()

906 + torch.save(audio_embedding, buffer)

907 + buffer.seek(0)

908 + binary_data = buffer.read()

909 + base64_audio_embedding = base64.b64encode(binary_data).decode("utf-8")

910 +

911 + conversation, mm_data, mm_uuids = parse_chat_messages(

912 + [

913 + {

914 + "role": "user",

915 + "content": [

916 + {"type": "text", "text": "Describe this audio"},

917 + {

918 + "type": "audio_embeds",

919 + "audio_embeds": base64_audio_embedding,

920 + },

921 + ],

922 + }

923 + ],

924 + audio_embeds_model_config,

925 + qwen2_audio_tokenizer,

926 + content_format="string",

927 + )

928 +

929 + # Should have audio embedding in mm_data (single tensor, not a list)

930 + assert mm_data is not None

931 + assert "audio" in mm_data

932 + assert isinstance(mm_data["audio"], torch.Tensor)

933 + assert mm_data["audio"].shape == audio_embedding.shape

934 + # No UUID provided

935 + assert mm_uuids is not None

936 + assert "audio" in mm_uuids

937 + _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[None])

938 +

939 +

940 +@pytest.mark.asyncio

941 +async def test_parse_chat_messages_audio_embeds_async(

942 + audio_embeds_model_config,

943 + qwen2_audio_tokenizer,

944 +):

945 + """Test audio_embeds with async futures."""

946 + import base64

947 + import io

948 +

949 + import torch

950 +

951 + # Create a sample audio embedding tensor

952 + audio_embedding = torch.randn(1, 128, 768)

953 +

954 + # Encode it as base64

955 + buffer = io.BytesIO()

956 + torch.save(audio_embedding, buffer)

957 + buffer.seek(0)

958 + binary_data = buffer.read()

959 + base64_audio_embedding = base64.b64encode(binary_data).decode("utf-8")

960 +

961 + conversation, mm_future, mm_uuids = parse_chat_messages_futures(

962 + [

963 + {

964 + "role": "user",

965 + "content": [

966 + {"type": "text", "text": "Describe this audio"},

967 + {

968 + "type": "audio_embeds",

969 + "audio_embeds": base64_audio_embedding,

970 + },

971 + ],

972 + }

973 + ],

974 + audio_embeds_model_config,

975 + qwen2_audio_tokenizer,

976 + content_format="string",

977 + )

978 +

979 + # Should have audio embedding in mm_data (single tensor, not a list)

980 + mm_data = await mm_future

981 + assert mm_data is not None

982 + assert "audio" in mm_data

983 + assert isinstance(mm_data["audio"], torch.Tensor)

984 + assert mm_data["audio"].shape == audio_embedding.shape

985 + # No UUID provided

986 + assert mm_uuids is not None

987 + assert "audio" in mm_uuids

988 + _assert_mm_uuids(mm_uuids, 1, modality="audio", expected_uuids=[None])

989 +

990 +

991 @pytest.mark.asyncio

992 async def test_parse_chat_messages_empty_image_embeds_with_uuid_async(

993 phi3v_model_config_image_embeds,

vllm/entrypoints/chat_utils.py

+147 -2

@@ -94,6 +94,22 @@ class ChatCompletionContentPartImageEmbedsParam(TypedDict, total=False):

94 """

97 +class ChatCompletionContentPartAudioEmbedsParam(TypedDict, total=False):

98 + audio_embeds: str | dict[str, str] | None

99 + """

100 + The audio embeddings. It can be either:

101 + - A single base64 string representing a serialized torch tensor.

102 + - A dictionary where each value is a base64 string.

103 + """

104 + type: Required[Literal["audio_embeds"]]

105 + """The type of the content part."""

106 + uuid: str | None

107 + """

108 + User-provided UUID of a media. User must guarantee that it is properly

109 + generated and unique for different medias.

110 + """

111 +

112 +

113 class VideoURL(TypedDict, total=False):

114 url: Required[str]

115 """

@@ -211,6 +227,7 @@ ChatCompletionContentPartParam: TypeAlias = (

227 | CustomChatCompletionContentPILImageParam

228 | CustomChatCompletionContentSimpleImageParam

229 | ChatCompletionContentPartImageEmbedsParam

230 + | ChatCompletionContentPartAudioEmbedsParam

231 | CustomChatCompletionContentSimpleAudioParam

232 | CustomChatCompletionContentSimpleVideoParam

233 | str

@@ -599,7 +616,7 @@ def resolve_chat_template_content_format(

616 return detected_format

617

618

602 -ModalityStr = Literal["image", "audio", "video", "image_embeds"]

619 +ModalityStr = Literal["image", "audio", "video", "image_embeds", "audio_embeds"]

620 _T = TypeVar("_T")

621

622

@@ -684,6 +701,11 @@ class BaseMultiModalItemTracker(ABC, Generic[_T]):

701 mm_uuids["image"] = uuids_by_modality["image_embeds"]

702 if "image" in uuids_by_modality:

703 mm_uuids["image"] = uuids_by_modality["image"] # UUIDs of images

704 + if "audio_embeds" in uuids_by_modality:

705 + audio_embeds_uuids = uuids_by_modality["audio_embeds"]

706 + if len(audio_embeds_uuids) > 1:

707 + raise ValueError("Only one message can have {'type': 'audio_embeds'}")

708 + mm_uuids["audio"] = uuids_by_modality["audio_embeds"]

709 if "audio" in uuids_by_modality:

710 mm_uuids["audio"] = uuids_by_modality["audio"] # UUIDs of audios

711 if "video" in uuids_by_modality:

@@ -703,6 +725,8 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[object]):

725 items_by_modality = dict(self._items_by_modality)

726 if "image" in items_by_modality and "image_embeds" in items_by_modality:

727 raise ValueError("Mixing raw image and embedding inputs is not allowed")

728 + if "audio" in items_by_modality and "audio_embeds" in items_by_modality:

729 + raise ValueError("Mixing raw audio and embedding inputs is not allowed")

730

731 if "image_embeds" in items_by_modality:

732 image_embeds_lst = items_by_modality["image_embeds"]

@@ -711,6 +735,11 @@ class MultiModalItemTracker(BaseMultiModalItemTracker[object]):

735 mm_inputs["image"] = image_embeds_lst[0]

736 if "image" in items_by_modality:

737 mm_inputs["image"] = items_by_modality["image"] # A list of images

738 + if "audio_embeds" in items_by_modality:

739 + audio_embeds_lst = items_by_modality["audio_embeds"]

740 + if len(audio_embeds_lst) > 1:

741 + raise ValueError("Only one message can have {'type': 'audio_embeds'}")

742 + mm_inputs["audio"] = audio_embeds_lst[0]

743 if "audio" in items_by_modality:

744 mm_inputs["audio"] = items_by_modality["audio"] # A list of audios

745 if "video" in items_by_modality:

@@ -738,6 +767,8 @@ class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):

767

768 if "image" in items_by_modality and "image_embeds" in items_by_modality:

769 raise ValueError("Mixing raw image and embedding inputs is not allowed")

770 + if "audio" in items_by_modality and "audio_embeds" in items_by_modality:

771 + raise ValueError("Mixing raw audio and embedding inputs is not allowed")

772

773 if "image_embeds" in items_by_modality:

774 image_embeds_lst = items_by_modality["image_embeds"]

@@ -746,6 +777,11 @@ class AsyncMultiModalItemTracker(BaseMultiModalItemTracker[Awaitable[object]]):

777 mm_inputs["image"] = image_embeds_lst[0]

778 if "image" in items_by_modality:

779 mm_inputs["image"] = items_by_modality["image"] # A list of images

780 + if "audio_embeds" in items_by_modality:

781 + audio_embeds_lst = items_by_modality["audio_embeds"]

782 + if len(audio_embeds_lst) > 1:

783 + raise ValueError("Only one message can have {'type': 'audio_embeds'}")

784 + mm_inputs["audio"] = audio_embeds_lst[0]

785 if "audio" in items_by_modality:

786 mm_inputs["audio"] = items_by_modality["audio"] # A list of audios

787 if "video" in items_by_modality:

@@ -804,6 +840,14 @@ class BaseMultiModalContentParser(ABC):

840 ) -> None:

841 raise NotImplementedError

842

843 + @abstractmethod

844 + def parse_audio_embeds(

845 + self,

846 + audio_embeds: str | dict[str, str] | None,

847 + uuid: str | None = None,

848 + ) -> None:

849 + raise NotImplementedError

850 +

851 @abstractmethod

852 def parse_video(self, video_url: str | None, uuid: str | None = None) -> None:

853 raise NotImplementedError

@@ -861,6 +905,31 @@ class MultiModalContentParser(BaseMultiModalContentParser):

905

906 self._add_placeholder("image", placeholder)

907

908 + def parse_audio_embeds(

909 + self,

910 + audio_embeds: str | dict[str, str] | None,

911 + uuid: str | None = None,

912 + ) -> None:

913 + mm_config = self.model_config.get_multimodal_config()

914 + if not mm_config.enable_mm_embeds:

915 + raise ValueError(

916 + "You must set `--enable-mm-embeds` to input `audio_embeds`"

917 + )

918 +

919 + if isinstance(audio_embeds, dict):

920 + embeds = {

921 + k: self._connector.fetch_audio_embedding(v)

922 + for k, v in audio_embeds.items()

923 + }

924 + placeholder = self._tracker.add("audio_embeds", embeds, uuid)

925 + elif isinstance(audio_embeds, str):

926 + embedding = self._connector.fetch_audio_embedding(audio_embeds)

927 + placeholder = self._tracker.add("audio_embeds", embedding, uuid)

928 + else:

929 + placeholder = self._tracker.add("audio_embeds", None, uuid)

930 +

931 + self._add_placeholder("audio", placeholder)

932 +

933 def parse_image_pil(

934 self, image_pil: Image.Image | None, uuid: str | None = None

935 ) -> None:

@@ -950,6 +1019,67 @@ class AsyncMultiModalContentParser(BaseMultiModalContentParser):

1019 placeholder = self._tracker.add("image_embeds", future, uuid)

1020 self._add_placeholder("image", placeholder)

1021

1022 + def parse_audio_embeds(

1023 + self,

1024 + audio_embeds: str | dict[str, str] | None,

1025 + uuid: str | None = None,

1026 + ) -> None:

1027 + mm_config = self.model_config.get_multimodal_config()

1028 + if not mm_config.enable_mm_embeds:

1029 + raise ValueError(

1030 + "You must set `--enable-mm-embeds` to input `audio_embeds`"

1031 + )

1032 +

1033 + logger.info(

1034 + "🎵 Parsing audio_embeds: type=%s, uuid=%s, is_dict=%s, "

1035 + "is_str=%s, is_none=%s",

1036 + type(audio_embeds).__name__,

1037 + uuid,

1038 + isinstance(audio_embeds, dict),

1039 + isinstance(audio_embeds, str),

1040 + audio_embeds is None,

1041 + )

1042 +

1043 + future: asyncio.Future[str | dict[str, str] | None] = asyncio.Future()

1044 +

1045 + if isinstance(audio_embeds, dict):

1046 + logger.info(

1047 + "🎵 Processing dict audio_embeds with %d entries",

1048 + len(audio_embeds),

1049 + )

1050 + embeds = {

1051 + k: self._connector.fetch_audio_embedding(v)

1052 + for k, v in audio_embeds.items()

1053 + }

1054 + future.set_result(embeds)

1055 + logger.info(

1056 + "🎵 Successfully loaded %d audio embeddings from dict",

1057 + len(embeds),

1058 + )

1059 +

1060 + if isinstance(audio_embeds, str):

1061 + base64_size = len(audio_embeds)

1062 + logger.info(

1063 + "🎵 Processing base64 audio_embeds: %d chars (%.2f KB)",

1064 + base64_size,

1065 + base64_size / 1024,

1066 + )

1067 + embedding = self._connector.fetch_audio_embedding(audio_embeds)

1068 + future.set_result(embedding)

1069 + logger.info(

1070 + "🎵 Successfully loaded audio embedding tensor: shape=%s, dtype=%s",

1071 + embedding.shape,

1072 + embedding.dtype,

1073 + )

1074 +

1075 + if audio_embeds is None:

1076 + logger.info("🎵 Audio embeds is None (UUID-only reference)")

1077 + future.set_result(None)

1078 +

1079 + placeholder = self._tracker.add("audio_embeds", future, uuid)

1080 + self._add_placeholder("audio", placeholder)

1081 + logger.info("🎵 Added audio_embeds placeholder with uuid=%s", uuid)

1082 +

1083 def parse_image_pil(

1084 self, image_pil: Image.Image | None, uuid: str | None = None

1085 ) -> None:

@@ -1132,6 +1262,7 @@ def _get_full_multimodal_text_prompt(

1262 # No need to validate using Pydantic again

1263 _TextParser = partial(cast, ChatCompletionContentPartTextParam)

1264 _ImageEmbedsParser = partial(cast, ChatCompletionContentPartImageEmbedsParam)

1265 +_AudioEmbedsParser = partial(cast, ChatCompletionContentPartAudioEmbedsParam)

1266 _InputAudioParser = partial(cast, ChatCompletionContentPartInputAudioParam)

1267 _RefusalParser = partial(cast, ChatCompletionContentPartRefusalParam)

1268 _PILImageParser = partial(cast, CustomChatCompletionContentPILImageParam)

@@ -1155,6 +1286,7 @@ MM_PARSER_MAP: dict[

1286 "input_image": lambda part: _ResponsesInputImageParser(part).get("image_url", None),

1287 "image_url": lambda part: _ImageParser(part).get("image_url", {}).get("url", None),

1288 "image_embeds": lambda part: _ImageEmbedsParser(part).get("image_embeds", None),

1289 + "audio_embeds": lambda part: _AudioEmbedsParser(part).get("audio_embeds", None),

1290 "image_pil": lambda part: _PILImageParser(part).get("image_pil", None),

1291 "audio_url": lambda part: _AudioParser(part).get("audio_url", {}).get("url", None),

1292 "input_audio": lambda part: _InputAudioParser(part).get("input_audio", None),

@@ -1223,8 +1355,17 @@ def _parse_chat_message_content_mm_part(

1355 )

1356 image_embeds = image_params.get("image_embeds", None)

1357 return "image_embeds", image_embeds

1358 + if "audio_embeds" in part:

1359 + # "audio_embeds" could be None if UUID is provided.

1360 + audio_params = cast( # type: ignore[assignment]

1361 + ChatCompletionContentPartAudioEmbedsParam, part

1362 + )

1363 + audio_embeds = audio_params.get("audio_embeds", None)

1364 + return "audio_embeds", audio_embeds

1365 if "audio_url" in part:

1227 - audio_params = cast(CustomChatCompletionContentSimpleAudioParam, part)

1366 + audio_params = cast( # type: ignore[assignment]

1367 + CustomChatCompletionContentSimpleAudioParam, part

1368 + )

1369 audio_url = audio_params.get("audio_url", None)

1370 if isinstance(audio_url, dict):

1371 # Can potentially happen if user provides a uuid

@@ -1348,6 +1489,10 @@ def _parse_chat_message_content_part(

1489 content = cast(str | dict[str, str], content) if content is not None else None

1490 mm_parser.parse_image_embeds(content, uuid)

1491 modality = "image"

1492 + elif part_type == "audio_embeds":

1493 + content = cast(str | dict[str, str], content) if content is not None else None

1494 + mm_parser.parse_audio_embeds(content, uuid)

1495 + modality = "audio"

1496 elif part_type == "audio_url":

1497 str_content = cast(str, content)

1498 mm_parser.parse_audio(str_content, uuid)

vllm/multimodal/audio.py

+24

@@ -7,6 +7,8 @@ from typing import Literal

8 import numpy as np

9 import numpy.typing as npt

10 +import pybase64

11 +import torch

13 from vllm.utils.import_utils import PlaceholderModule

@@ -116,3 +118,25 @@ class AudioMediaIO(MediaIO[tuple[npt.NDArray, float]]):

118 data = buffer.getvalue()

119

120 return base64.b64encode(data).decode("utf-8")

121 +

122 +

123 +class AudioEmbeddingMediaIO(MediaIO[torch.Tensor]):

124 + def __init__(self) -> None:

125 + super().__init__()

126 +

127 + def load_bytes(self, data: bytes) -> torch.Tensor:

128 + buffer = BytesIO(data)

129 + return torch.load(buffer, weights_only=True)

130 +

131 + def load_base64(self, media_type: str, data: str) -> torch.Tensor:

132 + return self.load_bytes(pybase64.b64decode(data, validate=True))

133 +

134 + def load_file(self, filepath: Path) -> torch.Tensor:

135 + return torch.load(filepath, weights_only=True)

136 +

137 + def encode_base64(self, media: torch.Tensor) -> str:

138 + buffer = BytesIO()

139 + torch.save(media, buffer)

140 + buffer.seek(0)

141 + binary_data = buffer.read()

142 + return pybase64.b64encode(binary_data).decode("utf-8")

vllm/multimodal/utils.py

+12 -1

@@ -22,7 +22,7 @@ from vllm.logger import init_logger

22 from vllm.utils.jsontree import json_map_leaves

23 from vllm.utils.registry import ExtensionManager

25 -from .audio import AudioMediaIO

25 +from .audio import AudioEmbeddingMediaIO, AudioMediaIO

26 from .base import MediaIO

27 from .image import ImageEmbeddingMediaIO, ImageMediaIO

28 from .video import VideoMediaIO

@@ -342,6 +342,17 @@ class MediaConnector:

342

343 return image_embedding_io.load_base64("", data)

344

345 + def fetch_audio_embedding(

346 + self,

347 + data: str,

348 + ) -> torch.Tensor:

349 + """

350 + Load audio embedding from a URL.

351 + """

352 + audio_embedding_io = AudioEmbeddingMediaIO()

353 +

354 + return audio_embedding_io.load_base64("", data)

355 +

356

357 def encode_audio_base64(

358 audio: np.ndarray,