[None][feat] Support audio inputs in OpenEngine gRPC

tensorrt_llm/_torch/models/modeling_whisper.py

@@ -598,6 +598,9 @@ class WhisperInputProcessor(InputProcessor):
     # Marks this model as feature-driven for the encoder side: prompts
     # without audio cannot be served and are rejected at submission.
     requires_encoder_features = True
+    # The text prompt selects the decoder language/task; pre-tokenized prompt
+    # IDs are not consumed by this processor.
+    supports_token_ids_input = False
 
     def __init__(self, model_path, config, tokenizer, trust_remote_code: bool = True, **kwargs):
         self.model_path = model_path

tensorrt_llm/grpc/openengine/README.md

@@ -69,6 +69,11 @@ Set `failureThreshold * periodSeconds` to cover the longest expected model start
 
 OpenEngine gRPC does not yet expose sleep or weight updates. When those operations are added, they must publish `inference=NOT_SERVING` for their duration while keeping `liveness=SERVING`; the current executor health predicate alone does not detect those admission states. A sidecar that loses `Watch` or sees a new engine identity must validate health again.
 
+For models that do not accept token-ID input, including Whisper, the optional
+`Health(include_inference_probe=True)` check reports the inference probe as not
+applicable. It does not submit inference and does not establish inference liveness;
+the ordinary lifecycle and executor health checks still apply.
+
 The `Inference.Generate` RPC loads the selected model through TensorRT-LLM's PyTorch `LLM` API and streams incremental token and text events. It supports text and token-ID inputs, native sampling and stopping options, top-N prompt and output log probabilities, TensorRT-LLM guided-decoding modes, cache salt, trace-context propagation, multiple output sequences, finish reasons, and final usage.
 
 Clients must continuously consume the response stream. If response delivery remains stalled for 30 seconds, the server aborts the engine request and terminates the stream with a retryable overload error.

tensorrt_llm/grpc/openengine/capabilities.py

@@ -64,30 +64,45 @@ def supported_guides(backend: str | None) -> frozenset[str]:
 
 def _loaded_model_type(llm: Any) -> str:
     """Return the top-level model type for the loaded input processor."""
-    return resolve_top_level_model_type(llm.input_processor.config)
+    input_processor = getattr(llm, "input_processor", None)
+    config = getattr(input_processor, "config", None) or getattr(llm, "_hf_model_config", None)
+    if config is None:
+        return ""
+    return resolve_top_level_model_type(config)
 
 
 def resolve_media_capabilities(
     llm: Any, implemented_modalities: Collection[str]
 ) -> MediaCapabilities:
     """Intersect model support with implemented request handlers.
 
-    A live multimodal processor prevents registry-only false positives.
+    A live registered processor prevents registry-only false positives.
     """
     input_processor = getattr(llm, "input_processor", None)
-    if not isinstance(input_processor, BaseMultimodalInputProcessor):
-        return MediaCapabilities({})
     if getattr(getattr(llm, "args", None), "disable_mm_encoder", False):
         return MediaCapabilities({})
+
     model_type = _loaded_model_type(llm)
     if not model_type:
         return MediaCapabilities({})
+    if not isinstance(input_processor, BaseMultimodalInputProcessor):
+        registered_model_type = getattr(type(input_processor), "_registered_model_type", None)
+        if registered_model_type != model_type:
+            return MediaCapabilities({})
+        if not callable(input_processor):
+            return MediaCapabilities({})
 
     placeholders = {
         modality: MULTIMODAL_PLACEHOLDER_REGISTRY.get_placeholder(model_type, modality)
         for modality in implemented_modalities
         if MULTIMODAL_PLACEHOLDER_REGISTRY.is_valid(model_type, modality)
     }
+    # Gemma4 checkpoints may omit the audio encoder despite sharing a model type.
+    if (
+        model_type == "gemma4"
+        and getattr(getattr(input_processor, "config", None), "audio_config", None) is None
+    ):
+        placeholders.pop("audio", None)
     # Fail startup before Control advertises a malformed template.
     for modality, template in placeholders.items():
         try:

tensorrt_llm/grpc/openengine/control.py

@@ -39,7 +39,7 @@
 from .coordinator import CoordinationError, FrontendClient, InvalidRequestIdError
 from .errors import AbortFailedError
 from .kv_events import data_parallel_size, events_config
-from .request_mapping import conversation_affinity_enabled
+from .request_mapping import conversation_affinity_enabled, token_ids_input_supported
 
 __all__ = ["OpenEngineControlServicer"]
 
@@ -289,7 +289,7 @@ async def GetModelInfo(
             served_model_name=self._model,
             # Without a tokenizer the input processor rejects every string prompt.
             supports_text_input=getattr(self._llm, "tokenizer", None) is not None,
-            supports_token_ids_input=True,
+            supports_token_ids_input=token_ids_input_supported(self._llm),
             # Generate rejects `lora_name` even on an enable_lora build.
             supports_lora=False,
             supports_multimodal=bool(media_modalities),
@@ -476,6 +476,15 @@ async def _inference_probe(self, model: str) -> lifecycle_pb2.HealthCheck:
     async def _run_inference_probe(self, model: str) -> lifecycle_pb2.HealthCheck:
         from tensorrt_llm.sampling_params import SamplingParams
 
+        if not token_ids_input_supported(self._llm):
+            return lifecycle_pb2.HealthCheck(
+                name="inference_probe",
+                state=lifecycle_pb2.HEALTH_STATE_READY,
+                message=(
+                    "probe not applicable: model does not accept token-ID input; "
+                    "inference liveness was not tested"
+                ),
+            )
         handle = None
         reservation = None
         # Counted by GetLoad and reachable by Abort(all_requests): a probe

tensorrt_llm/grpc/openengine/request_mapping.py

@@ -16,11 +16,13 @@
 
 from tensorrt_llm.conversation_params import ConversationParams
 from tensorrt_llm.inputs.media_io import (
+    AudioMediaIO,
     ImageMediaIO,
     MediaFetchError,
     VideoMediaIO,
     _decoded_video_bytes,
 )
+from tensorrt_llm.inputs.multimodal_data import AudioData
 from tensorrt_llm.inputs.registry import BaseMultimodalInputProcessor
 from tensorrt_llm.sampling_params import GuidedDecodingParams, SamplingParams
 
@@ -182,6 +184,12 @@ def conversation_affinity_enabled(llm: Any) -> bool:
     )
 
 
+def token_ids_input_supported(llm: Any) -> bool:
+    """Whether the loaded input processor consumes pre-tokenized prompts."""
+    input_processor = getattr(llm, "input_processor", None)
+    return getattr(input_processor, "supports_token_ids_input", True) is not False
+
+
 def _top_n_candidates(selection: Any, name: str) -> int:
     kind = selection.WhichOneof("selection")
     if kind is None:
@@ -395,12 +403,16 @@ def _metadata_from_context(
 
 
 def _base_input_from_request(
-    request: generation_pb2.GenerateRequest,
+    request: generation_pb2.GenerateRequest, llm: Any
 ) -> str | dict[str, list[int]]:
     input_kind = request.WhichOneof("input")
     if input_kind == "prompt":
         return request.prompt
     if input_kind == "token_ids":
+        if not token_ids_input_supported(llm):
+            raise UnsupportedFeatureError(
+                "token-ID input is not supported by the loaded model; send a text prompt"
+            )
         return {"prompt_token_ids": request.token_ids.ids[:]}
     raise ValueError("exactly one of prompt or token_ids must be set")
 
@@ -459,6 +471,10 @@ def _validate_image_item(item: generation_pb2.MediaItem, index: int) -> None:
     _validate_media_item(item, index, "image")
 
 
+def _validate_audio_item(item: generation_pb2.MediaItem, index: int) -> None:
+    _validate_media_item(item, index, "audio")
+
+
 def _validate_video_item(item: generation_pb2.MediaItem, index: int) -> None:
     _validate_media_item(item, index, "video")
 
@@ -507,6 +523,17 @@ async def load(item: generation_pb2.MediaItem, index: int) -> Any:
     return load
 
 
+def _audio_loader_factory(llm: Any, budget: _DecodedMediaBudget) -> _MediaLoader:
+    audio_io = _media_io_for_model(
+        llm, "audio", AudioMediaIO, {"decoded_memory_reserver": budget.reserve}
+    )
+
+    async def load(item: generation_pb2.MediaItem, index: int) -> Any:
+        return await _load_media_item(item, audio_io, index, "audio")
+
+    return load
+
+
 def _decoded_image_bytes(image: Any) -> int:
     import numpy as np
     import torch
@@ -521,6 +548,23 @@ def _decoded_image_bytes(image: Any) -> int:
     raise TypeError(f"unsupported decoded image type: {type(image).__name__}")
 
 
+def _decoded_audio_data_bytes(audio: AudioData) -> int:
+    import numpy as np
+    import torch
+
+    if isinstance(audio.samples, torch.Tensor):
+        return audio.samples.numel() * audio.samples.element_size()
+    if isinstance(audio.samples, np.ndarray):
+        return audio.samples.nbytes
+    raise TypeError(f"unsupported decoded audio type: {type(audio.samples).__name__}")
+
+
+def _decoded_audio_bytes(audio: Any) -> int:
+    if not isinstance(audio, tuple) or len(audio) != 2:
+        raise TypeError(f"unsupported decoded audio type: {type(audio).__name__}")
+    return _decoded_audio_data_bytes(AudioData(samples=audio[0], sample_rate=audio[1]))
+
+
 def _video_runtime_available() -> bool:
     try:
         importlib.import_module("cv2")
@@ -541,6 +585,13 @@ def _video_runtime_available() -> bool:
         _video_runtime_available,
         reserves_decoded_bytes=True,
     ),
+    generation_pb2.MODALITY_AUDIO: _MediaHandler(
+        "audio",
+        _validate_audio_item,
+        _audio_loader_factory,
+        _decoded_audio_bytes,
+        reserves_decoded_bytes=True,
+    ),
 }
 
 
@@ -564,7 +615,7 @@ def _prepare_input_from_request(
     media_capabilities: MediaCapabilities | None = None,
 ) -> _PreparedInput:
     """Validate request media without starting any fetch or decode work."""
-    inputs = _base_input_from_request(request)
+    inputs = _base_input_from_request(request, llm)
     if not request.media:
         return _PreparedInput(inputs, (), None)
 
@@ -630,6 +681,8 @@ def _validate_media_placeholders(prepared: _PreparedInput, llm: Any) -> None:
     if not prepared.media or prepared.capabilities is None:
         return
     processor = llm.input_processor
+    if not isinstance(processor, BaseMultimodalInputProcessor):
+        return
     if not processor.requires_input_placeholders(tokenized=not isinstance(prepared.inputs, str)):
         return
     tokenizer = llm.tokenizer
@@ -750,4 +803,5 @@ async def _input_from_request(
 __all__ = [
     "conversation_params_from_request",
     "sampling_params_from_request",
+    "token_ids_input_supported",
 ]

tensorrt_llm/grpc/openengine/server.py

@@ -23,6 +23,7 @@
 
 from tensorrt_llm import LLM as PyTorchLLM
 from tensorrt_llm.grpc._protobuf import _warn_if_python_protobuf
+from tensorrt_llm.inputs.media_io import _close_aiohttp_session
 from tensorrt_llm.logger import logger
 from tensorrt_llm.serve._frontend_processes import (
     FrontendStartupCancelled,
@@ -276,10 +277,17 @@ async def stop(self, grace: float = 5.0) -> None:
                 await asyncio.gather(self._health_task, return_exceptions=True)
             await self._server.stop(grace=grace)
         finally:
-            await asyncio.gather(
-                asyncio.to_thread(self._input_executor.shutdown, wait=True, cancel_futures=True),
-                asyncio.to_thread(self._media_executor.shutdown, wait=True, cancel_futures=True),
-            )
+            try:
+                await _close_aiohttp_session()
+            finally:
+                await asyncio.gather(
+                    asyncio.to_thread(
+                        self._input_executor.shutdown, wait=True, cancel_futures=True
+                    ),
+                    asyncio.to_thread(
+                        self._media_executor.shutdown, wait=True, cancel_futures=True
+                    ),
+                )
         logger.info("OpenEngine server stopped")
 
     async def wait_for_termination(self) -> None: