feat(multimodal): forward mm_processor_kwargs through pipeline for use_audio_in_video (#8150)

Co-authored-by: Claude Opus 4.6 (1M context) <noreply@anthropic.com>

components/src/dynamo/frontend/vllm_processor.py

@@ -277,6 +277,12 @@ class VllmProcessor:
         if mm_data:
             dynamo_preproc["multi_modal_data"] = mm_data
 
+        # Forward mm_processor_kwargs (e.g. use_audio_in_video) to the backend.
+        if request_for_sampling.mm_processor_kwargs is not None:
+            dynamo_preproc[
+                "mm_processor_kwargs"
+            ] = request_for_sampling.mm_processor_kwargs
+
         post = StreamingPostProcessor(
             tokenizer=self.tokenizer,
             request_for_sampling=request_for_sampling,
@@ -310,6 +316,10 @@ class VllmProcessor:
 
         try:
             if self.is_kv_router:
+                extra_args: dict[str, Any] = {}
+                mm_proc_kwargs = dynamo_preproc.get("mm_processor_kwargs")
+                if mm_proc_kwargs is not None:
+                    extra_args["mm_processor_kwargs"] = mm_proc_kwargs
                 dynamo_stream = await self.router.generate(
                     token_ids=tokens,
                     model=dynamo_preproc["model"],
@@ -317,6 +327,7 @@ class VllmProcessor:
                     sampling_options=dynamo_preproc["sampling_options"],
                     output_options=dynamo_preproc["output_options"],
                     multi_modal_data=dynamo_preproc.get("multi_modal_data"),
+                    extra_args=extra_args or None,
                 )
             else:
                 dynamo_stream = await self.router.generate(

components/src/dynamo/router/__main__.py

@@ -113,6 +113,7 @@ class StandaloneRouterHandler:
             "prefill_result": request.get("prefill_result"),
             "bootstrap_info": request.get("bootstrap_info"),
             "extra_args": request.get("extra_args"),
+            "mm_processor_kwargs": request.get("mm_processor_kwargs"),
         }
 
         async for worker_output in await self.kv_router.generate_from_request(

components/src/dynamo/vllm/handlers.py

@@ -1179,8 +1179,28 @@ class BaseWorkerHandler(ABC, Generic[RequestT, ResponseT]):
 
         return prompt, sequence_length, embeddings_tensor
 
+    @staticmethod
+    def _get_mm_processor_kwargs(
+        request: Dict[str, Any],
+    ) -> Dict[str, Any] | None:
+        """Extract mm_processor_kwargs from a request dict.
+
+        Checks the top-level key (client router / Rust preprocessor path)
+        and falls back to ``extra_args`` (KV router path).
+        """
+        mm_processor_kwargs = request.get("mm_processor_kwargs")
+        if mm_processor_kwargs is None:
+            req_extra_args = request.get("extra_args")
+            if isinstance(req_extra_args, dict):
+                mm_processor_kwargs = req_extra_args.get("mm_processor_kwargs")
+        return mm_processor_kwargs
+
     async def _extract_multimodal_data(
-        self, request: Dict[str, Any], request_id: str, context
+        self,
+        request: Dict[str, Any],
+        request_id: str,
+        context,
+        mm_processor_kwargs: Dict[str, Any] | None = None,
     ) -> Dict[str, Any] | None:
         """
         Extract and decode multimodal data from PreprocessedRequest.
@@ -1254,6 +1274,54 @@ class BaseWorkerHandler(ABC, Generic[RequestT, ResponseT]):
                     f"Extracted {len(audios)} audio item(s) for multimodal processing"
                 )
 
+        # Extract audio from video URLs when use_audio_in_video is set.
+        # Models expect 1:1 audio/video pairing in the same order.
+        # We load per-video sequentially to preserve ordering; a video
+        # without an audio track raises immediately to avoid corrupting
+        # the alignment.
+        if (
+            video_mm_items
+            and mm_processor_kwargs
+            and mm_processor_kwargs.get("use_audio_in_video", False)
+        ):
+            video_audios: list = []
+            for item in video_mm_items:
+                url = item.get(URL_VARIANT_KEY) if isinstance(item, dict) else None
+                if not url:
+                    raise ValueError(
+                        "use_audio_in_video requires all video items to be "
+                        "URL-based. Got a non-URL video item (e.g. frontend-"
+                        "decoded). Audio extraction from decoded video data "
+                        "is not yet supported."
+                    )
+                try:
+                    audio = await self.audio_loader.load_audio(url)
+                    video_audios.append(audio)
+                except Exception:
+                    logger.error(
+                        "Failed to extract audio from video %s. "
+                        "use_audio_in_video requires every video to "
+                        "contain an audio stream.",
+                        url[:80],
+                    )
+                    raise
+            if video_audios:
+                existing = vllm_mm_data.get("audio")
+                if existing is not None:
+                    all_audios = (
+                        existing if isinstance(existing, list) else [existing]
+                    ) + video_audios
+                else:
+                    all_audios = video_audios
+                vllm_mm_data["audio"] = (
+                    all_audios[0] if len(all_audios) == 1 else all_audios
+                )
+                logger.debug(
+                    "Extracted %d audio track(s) from video URL(s) "
+                    "(use_audio_in_video=True)",
+                    len(video_audios),
+                )
+
         return vllm_mm_data if vllm_mm_data else None
 
     def _build_prompt_from_request(
@@ -1262,6 +1330,7 @@ class BaseWorkerHandler(ABC, Generic[RequestT, ResponseT]):
         request_id: str,
         multi_modal_data: Dict[str, Any] | None,
         log_prefix: str = "",
+        mm_processor_kwargs: Dict[str, Any] | None = None,
     ) -> tuple[TokensPrompt | EmbedsPrompt | None, int | None, Dict[str, Any] | None]:
         """
         Build a prompt from request, handling both prompt_embeds and token_ids.
@@ -1271,6 +1340,8 @@ class BaseWorkerHandler(ABC, Generic[RequestT, ResponseT]):
             request_id: Request ID for logging
             multi_modal_data: Optional multimodal data to attach to TokensPrompt
             log_prefix: Prefix for log messages (e.g., "Prefill " for prefill requests)
+            mm_processor_kwargs: Optional multimodal processor kwargs (e.g.
+                use_audio_in_video) forwarded to the vLLM engine.
 
         Returns:
             Tuple of (prompt, embedding_sequence_length, error_dict) where:
@@ -1313,6 +1384,8 @@ class BaseWorkerHandler(ABC, Generic[RequestT, ResponseT]):
         )
         if mm_uuids is not None:
             prompt_kwargs["multi_modal_uuids"] = mm_uuids
+        if mm_processor_kwargs is not None:
+            prompt_kwargs["mm_processor_kwargs"] = mm_processor_kwargs
 
         prompt = TokensPrompt(**prompt_kwargs)
         return prompt, embedding_sequence_length, None
@@ -1616,6 +1689,8 @@ class DecodeWorkerHandler(BaseWorkerHandler):
             "multi_modal_data" in request and request["multi_modal_data"] is not None
         )
 
+        mm_processor_kwargs = self._get_mm_processor_kwargs(request)
+
         multi_modal_data = None
         if is_decode_only:
             # Decode mode: branch on model, not data.
@@ -1650,17 +1725,26 @@ class DecodeWorkerHandler(BaseWorkerHandler):
                 mm = request["multi_modal_data"]
                 if mm.get(VIDEO_URL_KEY) or mm.get(AUDIO_URL_KEY):
                     multi_modal_data = await self._extract_multimodal_data(
-                        request, request_id, context
+                        request,
+                        request_id,
+                        context,
+                        mm_processor_kwargs=mm_processor_kwargs,
                     )
         else:
             # Aggregated mode: load images normally
             multi_modal_data = await self._extract_multimodal_data(
-                request, request_id, context
+                request,
+                request_id,
+                context,
+                mm_processor_kwargs=mm_processor_kwargs,
             )
 
         # Build prompt from request (handles both prompt_embeds and token_ids)
         prompt, embedding_sequence_length, error = self._build_prompt_from_request(
-            request, request_id, multi_modal_data
+            request,
+            request_id,
+            multi_modal_data,
+            mm_processor_kwargs=mm_processor_kwargs,
         )
         if error is not None:
             yield error
@@ -1867,15 +1951,24 @@ class PrefillWorkerHandler(BaseWorkerHandler):
 
     async def _generate_token_mode(self, request, context, request_id):
         """Generate prefill using internal protocol format (token-in-token-out)."""
+        mm_processor_kwargs = self._get_mm_processor_kwargs(request)
+
         # Extract and decode multimodal data if present
         multi_modal_data = await self._extract_multimodal_data(
-            request, request_id, context
+            request,
+            request_id,
+            context,
+            mm_processor_kwargs=mm_processor_kwargs,
         )
         embedding_params = self._build_embedding_params(multi_modal_data or {})
 
         # Build prompt from request (handles both prompt_embeds and token_ids)
         prompt, embedding_sequence_length, error = self._build_prompt_from_request(
-            request, request_id, multi_modal_data, log_prefix="Prefill "
+            request,
+            request_id,
+            multi_modal_data,
+            log_prefix="Prefill ",
+            mm_processor_kwargs=mm_processor_kwargs,
         )
         if error is not None:
             # Prefill errors need disaggregated_params field

components/src/dynamo/vllm/tests/test_vllm_video_handler.py

@@ -29,6 +29,10 @@ def _make_handler(enable_multimodal: bool = True) -> _TestWorkerHandler:
     handler.embedding_loader = None
     handler.image_loader = SimpleNamespace(load_image_batch=AsyncMock(return_value=[]))
     handler.video_loader = SimpleNamespace(load_video_batch=AsyncMock(return_value=[]))
+    handler.audio_loader = SimpleNamespace(
+        load_audio_batch=AsyncMock(return_value=[]),
+        load_audio=AsyncMock(return_value=(np.zeros(16000, dtype=np.float32), 16000.0)),
+    )
     return handler
 
 
@@ -128,3 +132,235 @@ async def test_extract_multimodal_data_rejects_requests_when_disabled():
             "req-4",
             context=None,
         )
+
+
+# --- use_audio_in_video tests ---
+
+
+@pytest.mark.asyncio
+async def test_extract_audio_from_video_when_use_audio_in_video():
+    """Audio is extracted from video URLs when use_audio_in_video=True."""
+    handler = _make_handler()
+    video = (
+        np.zeros((2, 4, 4, 3), dtype=np.uint8),
+        {"fps": 2.0, "frames_indices": [0, 1], "total_num_frames": 2},
+    )
+    audio = (np.zeros(16000, dtype=np.float32), 16000.0)
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video])
+    handler.audio_loader.load_audio = AsyncMock(return_value=audio)
+
+    result = await handler._extract_multimodal_data(
+        {"multi_modal_data": {"video_url": [{"Url": "https://example.com/video.mp4"}]}},
+        "req-aiv-1",
+        context=None,
+        mm_processor_kwargs={"use_audio_in_video": True},
+    )
+
+    assert result is not None
+    assert result["video"] is video
+    assert result["audio"] is audio
+    handler.audio_loader.load_audio.assert_awaited_once_with(
+        "https://example.com/video.mp4"
+    )
+
+
+@pytest.mark.asyncio
+async def test_no_audio_from_video_without_flag():
+    """Without use_audio_in_video, no audio is extracted from video URLs."""
+    handler = _make_handler()
+    video = (
+        np.zeros((2, 4, 4, 3), dtype=np.uint8),
+        {"fps": 2.0, "frames_indices": [0, 1], "total_num_frames": 2},
+    )
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video])
+
+    result = await handler._extract_multimodal_data(
+        {"multi_modal_data": {"video_url": [{"Url": "https://example.com/video.mp4"}]}},
+        "req-aiv-2",
+        context=None,
+    )
+
+    assert result is not None
+    assert result["video"] is video
+    assert "audio" not in result
+    handler.audio_loader.load_audio.assert_not_awaited()
+
+
+@pytest.mark.asyncio
+async def test_audio_from_video_multiple_videos_preserves_order():
+    """With multiple videos, audio is extracted per-video in the same order."""
+    handler = _make_handler()
+    video_a = (np.zeros((1, 4, 4, 3), dtype=np.uint8), {"fps": 1.0})
+    video_b = (np.ones((1, 4, 4, 3), dtype=np.uint8), {"fps": 1.0})
+    audio_a = (np.zeros(8000, dtype=np.float32), 16000.0)
+    audio_b = (np.ones(8000, dtype=np.float32), 16000.0)
+
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video_a, video_b])
+    handler.audio_loader.load_audio = AsyncMock(side_effect=[audio_a, audio_b])
+
+    video_items = [
+        {"Url": "https://example.com/a.mp4"},
+        {"Url": "https://example.com/b.mp4"},
+    ]
+    result = await handler._extract_multimodal_data(
+        {"multi_modal_data": {"video_url": video_items}},
+        "req-aiv-3",
+        context=None,
+        mm_processor_kwargs={"use_audio_in_video": True},
+    )
+
+    assert result is not None
+    assert result["video"] == [video_a, video_b]
+    assert result["audio"] == [audio_a, audio_b]
+
+
+@pytest.mark.asyncio
+async def test_audio_from_video_raises_on_silent_video():
+    """A video without an audio track raises — silent videos break 1:1 ordering."""
+    handler = _make_handler()
+    video = (
+        np.zeros((2, 4, 4, 3), dtype=np.uint8),
+        {"fps": 2.0, "frames_indices": [0, 1], "total_num_frames": 2},
+    )
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video])
+    handler.audio_loader.load_audio = AsyncMock(
+        side_effect=RuntimeError("no audio stream")
+    )
+
+    with pytest.raises(RuntimeError, match="no audio stream"):
+        await handler._extract_multimodal_data(
+            {
+                "multi_modal_data": {
+                    "video_url": [{"Url": "https://example.com/silent.mp4"}]
+                }
+            },
+            "req-aiv-4",
+            context=None,
+            mm_processor_kwargs={"use_audio_in_video": True},
+        )
+
+
+@pytest.mark.asyncio
+async def test_audio_from_video_raises_on_non_url_video():
+    """A decoded (non-URL) video item raises when use_audio_in_video is set."""
+    handler = _make_handler()
+    video = (np.zeros((2, 4, 4, 3), dtype=np.uint8), {"fps": 2.0})
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video])
+
+    with pytest.raises(ValueError, match="non-URL video item"):
+        await handler._extract_multimodal_data(
+            {"multi_modal_data": {"video_url": [{"Decoded": {"shape": [2, 4, 4, 3]}}]}},
+            "req-aiv-decoded",
+            context=None,
+            mm_processor_kwargs={"use_audio_in_video": True},
+        )
+
+
+@pytest.mark.asyncio
+async def test_audio_from_video_merges_with_standalone_audio():
+    """Standalone audio_url items and video-extracted audio are both included."""
+    handler = _make_handler()
+    video = (np.zeros((2, 4, 4, 3), dtype=np.uint8), {"fps": 2.0})
+    standalone_audio = (np.zeros(8000, dtype=np.float32), 16000.0)
+    video_audio = (np.ones(8000, dtype=np.float32), 16000.0)
+
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video])
+    handler.audio_loader.load_audio_batch = AsyncMock(return_value=[standalone_audio])
+    handler.audio_loader.load_audio = AsyncMock(return_value=video_audio)
+
+    result = await handler._extract_multimodal_data(
+        {
+            "multi_modal_data": {
+                "video_url": [{"Url": "https://example.com/video.mp4"}],
+                "audio_url": [{"Url": "https://example.com/narration.wav"}],
+            }
+        },
+        "req-aiv-5",
+        context=None,
+        mm_processor_kwargs={"use_audio_in_video": True},
+    )
+
+    assert result is not None
+    assert result["audio"] == [standalone_audio, video_audio]
+
+
+@pytest.mark.asyncio
+async def test_build_prompt_includes_mm_processor_kwargs():
+    """mm_processor_kwargs is included in the TokensPrompt."""
+    handler = _make_handler()
+    mm_kwargs = {"use_audio_in_video": True}
+
+    prompt, _, error = handler._build_prompt_from_request(
+        {"token_ids": [1, 2, 3]},
+        "req-prompt-1",
+        multi_modal_data=None,
+        mm_processor_kwargs=mm_kwargs,
+    )
+
+    assert error is None
+    assert prompt["mm_processor_kwargs"] is mm_kwargs
+
+
+@pytest.mark.asyncio
+async def test_build_prompt_excludes_mm_processor_kwargs_when_none():
+    """mm_processor_kwargs is not added to TokensPrompt when None."""
+    handler = _make_handler()
+
+    prompt, _, error = handler._build_prompt_from_request(
+        {"token_ids": [1, 2, 3]},
+        "req-prompt-2",
+        multi_modal_data=None,
+    )
+
+    assert error is None
+    assert "mm_processor_kwargs" not in prompt
+
+
+# --- extra_args extraction tests (KV router path) ---
+
+
+@pytest.mark.asyncio
+async def test_extract_audio_from_video_with_mm_kwargs_in_extra_args():
+    """mm_processor_kwargs nested in extra_args (KV router path) triggers
+    audio extraction the same way as top-level mm_processor_kwargs."""
+    handler = _make_handler()
+    video = (
+        np.zeros((2, 4, 4, 3), dtype=np.uint8),
+        {"fps": 2.0, "frames_indices": [0, 1], "total_num_frames": 2},
+    )
+    audio = (np.zeros(16000, dtype=np.float32), 16000.0)
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video])
+    handler.audio_loader.load_audio = AsyncMock(return_value=audio)
+
+    result = await handler._extract_multimodal_data(
+        {"multi_modal_data": {"video_url": [{"Url": "https://example.com/video.mp4"}]}},
+        "req-extra-1",
+        context=None,
+        mm_processor_kwargs={"use_audio_in_video": True},
+    )
+
+    assert result is not None
+    assert result["audio"] is audio
+
+
+@pytest.mark.asyncio
+async def test_no_audio_extraction_when_extra_args_lacks_mm_kwargs():
+    """Without mm_processor_kwargs (neither top-level nor in extra_args),
+    no audio is extracted from video URLs."""
+    handler = _make_handler()
+    video = (
+        np.zeros((2, 4, 4, 3), dtype=np.uint8),
+        {"fps": 2.0, "frames_indices": [0, 1], "total_num_frames": 2},
+    )
+    handler.video_loader.load_video_batch = AsyncMock(return_value=[video])
+
+    result = await handler._extract_multimodal_data(
+        {"multi_modal_data": {"video_url": [{"Url": "https://example.com/video.mp4"}]}},
+        "req-extra-2",
+        context=None,
+        mm_processor_kwargs=None,
+    )
+
+    assert result is not None
+    assert "audio" not in result
+    handler.audio_loader.load_audio.assert_not_awaited()

lib/llm/src/preprocessor.rs

@@ -344,6 +344,9 @@ impl OpenAIPreprocessor {
             }));
         }
 
+        // Forward mm_processor_kwargs (e.g. use_audio_in_video) to the backend.
+        builder.mm_processor_kwargs(request.mm_processor_kwargs().cloned());
+
         Ok(builder)
     }
 

lib/llm/src/preprocessor/prompt.rs

@@ -89,6 +89,10 @@ pub trait OAIChatLikeRequest {
     fn media_io_kwargs(&self) -> Option<&MediaDecoder> {
         None
     }
+
+    fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> {
+        None
+    }
 }
 
 pub trait OAIPromptFormatter: Send + Sync + 'static {

lib/llm/src/preprocessor/prompt/template/oai.rs

@@ -346,6 +346,10 @@ impl OAIChatLikeRequest for NvCreateChatCompletionRequest {
     fn media_io_kwargs(&self) -> Option<&MediaDecoder> {
         self.media_io_kwargs.as_ref()
     }
+
+    fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> {
+        self.inner.mm_processor_kwargs.as_ref()
+    }
 }
 
 impl OAIChatLikeRequest for NvCreateCompletionRequest {

lib/llm/src/protocols/common/preprocessor.rs

@@ -195,6 +195,12 @@ pub struct PreprocessedRequest {
     #[serde(default, skip_serializing_if = "Option::is_none")]
     pub extra_args: Option<serde_json::Value>,
 
+    /// Multimodal processor kwargs forwarded to the backend engine
+    /// (e.g. `{"use_audio_in_video": true}` for omni models).
+    #[builder(default)]
+    #[serde(default, skip_serializing_if = "Option::is_none")]
+    pub mm_processor_kwargs: Option<serde_json::Value>,
+
     /// Optional request timestamp in milliseconds forwarded from nvext.
     #[builder(default)]
     #[serde(default, skip_serializing_if = "Option::is_none")]

lib/llm/src/protocols/unified.rs

@@ -485,6 +485,10 @@ impl OAIChatLikeRequest for UnifiedRequest {
     fn media_io_kwargs(&self) -> Option<&MediaDecoder> {
         self.inner.media_io_kwargs.as_ref()
     }
+
+    fn mm_processor_kwargs(&self) -> Option<&serde_json::Value> {
+        self.inner.inner.mm_processor_kwargs.as_ref()
+    }
 }
 
 impl UnifiedRequest {