refactor(grpc): centralize backend×modality capability check

model_gateway/src/routers/grpc/client.rs

@@ -184,6 +184,19 @@ impl GrpcClient {
         matches!(self, Self::TokenSpeed(_))
     }
 
+    /// Runtime type backing this client. Lets shared logic (e.g. the multimodal
+    /// capability matrix) key on the backend without matching every variant.
+    pub fn runtime_type(&self) -> crate::worker::RuntimeType {
+        use crate::worker::RuntimeType;
+        match self {
+            Self::Sglang(_) => RuntimeType::Sglang,
+            Self::Vllm(_) => RuntimeType::Vllm,
+            Self::Trtllm(_) => RuntimeType::Trtllm,
+            Self::Mlx(_) => RuntimeType::Mlx,
+            Self::TokenSpeed(_) => RuntimeType::TokenSpeed,
+        }
+    }
+
     pub async fn connect(
         url: &str,
         runtime_type: &str,

model_gateway/src/routers/grpc/common/stages/worker_selection.rs

@@ -170,6 +170,23 @@ impl PipelineStage for WorkerSelectionStage {
             }
         };
 
+        // Reject an unsupported (backend, modality) combination now that the
+        // runtime is known, before request building fetches/preprocesses media
+        // only to fail deep in assembly. The prefill leg builds the request in
+        // disaggregated mode, so its runtime is the one that must support the
+        // request's modalities.
+        if let Some(intermediate) = multimodal_intermediate(prep) {
+            if let Err(err) = multimodal::ensure_backend_supports_modalities(
+                selection_runtime(&workers),
+                intermediate,
+            ) {
+                return Err(error::bad_request(
+                    "multimodal_not_supported",
+                    format!("{err}"),
+                ));
+            }
+        }
+
         ctx.state.workers = Some(workers);
         Ok(None)
     }
@@ -179,6 +196,30 @@ impl PipelineStage for WorkerSelectionStage {
     }
 }
 
+/// Runtime of the leg that builds the generate request: the sole worker in
+/// regular mode, the prefill worker in disaggregated (PD/EPD) mode.
+fn selection_runtime(workers: &WorkerSelection) -> RuntimeType {
+    match workers {
+        WorkerSelection::Single { worker } => worker.metadata().spec.runtime_type,
+        WorkerSelection::Disaggregated { runtime_type, .. } => *runtime_type,
+    }
+}
+
+/// Borrow the request's multimodal intermediate, if any.
+fn multimodal_intermediate(
+    prep: &PreparationOutput,
+) -> Option<&multimodal::MultimodalIntermediate> {
+    match prep {
+        PreparationOutput::Chat {
+            processed_messages, ..
+        }
+        | PreparationOutput::Messages {
+            processed_messages, ..
+        } => processed_messages.multimodal_intermediate.as_ref(),
+        _ => None,
+    }
+}
+
 impl WorkerSelectionStage {
     fn select_single_worker(
         &self,
@@ -558,16 +599,7 @@ impl WorkerSelectionStage {
 }
 
 fn encode_item_hashes(prep: &PreparationOutput) -> anyhow::Result<Vec<Vec<u8>>> {
-    let intermediate = match prep {
-        PreparationOutput::Chat {
-            processed_messages, ..
-        }
-        | PreparationOutput::Messages {
-            processed_messages, ..
-        } => processed_messages.multimodal_intermediate.as_ref(),
-        _ => None,
-    };
-    let Some(intermediate) = intermediate else {
+    let Some(intermediate) = multimodal_intermediate(prep) else {
         return Ok(Vec::new());
     };
     multimodal::encode_routing_hashes(intermediate)

model_gateway/src/routers/grpc/multimodal/assemble.rs

@@ -19,6 +19,7 @@ use smg_grpc_client::common_proto as common;
 use tracing::{info, warn};
 
 use super::{
+    capability::ensure_backend_supports_modalities,
     log_mm_timing_enabled,
     serialize::{
         model_specific_to_tensor_bytes, serialize_array_as_tokenspeed_tensor,
@@ -66,17 +67,21 @@ async fn assemble_multimodal_data_impl(
     omit_prefill_pixels: bool,
 ) -> Result<MultimodalData> {
     validate_intermediate(&intermediate)?;
+    // Defense in depth: worker selection already rejected unsupported (backend,
+    // modality) combinations. Re-assert here against the same capability matrix
+    // so assembly can never silently mishandle a modality the backend lacks.
+    ensure_client_supports_intermediate(client, &intermediate)?;
     match client {
         GrpcClient::Sglang(_) => {
-            let batch = into_single_image_batch(intermediate, "SGLang")?;
+            let batch = into_single_batch(intermediate, "SGLang")?;
             Ok(MultimodalData::Sglang(assemble_sglang(batch)?))
         }
         GrpcClient::Vllm(_) => {
-            let batch = into_single_vision_batch(intermediate, "vLLM")?;
+            let batch = into_single_batch(intermediate, "vLLM")?;
             Ok(MultimodalData::Vllm(assemble_vllm(batch, workers)?))
         }
         GrpcClient::Trtllm(_) => {
-            let batch = into_single_image_batch(intermediate, "TRT-LLM")?;
+            let batch = into_single_batch(intermediate, "TRT-LLM")?;
             Ok(MultimodalData::Trtllm(assemble_trtllm(batch)?))
         }
         GrpcClient::TokenSpeed(_) => {
@@ -132,46 +137,34 @@ impl Drop for PendingTokenSpeedAssembly {
     }
 }
 
-fn into_single_image_batch(
+/// Backends other than TokenSpeed take a single preprocessed batch (one
+/// modality). The per-modality capability is enforced by
+/// [`ensure_client_supports_intermediate`]; this only enforces the structural
+/// single-batch constraint (multiple modalities in one request are TokenSpeed-
+/// only).
+fn into_single_batch(
     intermediate: MultimodalIntermediate,
     backend: &str,
 ) -> Result<PrecomputedMultimodalIntermediate> {
     anyhow::ensure!(
         intermediate.batches().len() == 1,
-        "{backend} multimodal path requires exactly one image batch; got {} batches",
+        "{backend} multimodal path requires exactly one batch; got {} batches",
         intermediate.batches().len()
     );
-    let mut batches = intermediate.into_batches();
-    let batch = batches
+    intermediate
+        .into_batches()
         .pop()
-        .context("multimodal intermediate is missing its sole batch")?;
-    anyhow::ensure!(
-        matches!(&batch.media, MediaBatch::Images(_)),
-        "{backend} multimodal path currently supports image inputs only; got {}",
-        batch.media.modality()
-    );
-    Ok(batch)
+        .context("multimodal intermediate is missing its sole batch")
 }
 
-fn into_single_vision_batch(
-    intermediate: MultimodalIntermediate,
-    backend: &str,
-) -> Result<PrecomputedMultimodalIntermediate> {
-    anyhow::ensure!(
-        intermediate.batches().len() == 1,
-        "{backend} multimodal path requires exactly one vision batch; got {} batches",
-        intermediate.batches().len()
-    );
-    let mut batches = intermediate.into_batches();
-    let batch = batches
-        .pop()
-        .context("multimodal intermediate is missing its sole batch")?;
-    anyhow::ensure!(
-        matches!(&batch.media, MediaBatch::Images(_) | MediaBatch::Videos(_)),
-        "{backend} multimodal path currently supports image and video inputs only; got {}",
-        batch.media.modality()
-    );
-    Ok(batch)
+/// Defense-in-depth capability assertion mirroring the early worker-selection
+/// check, keyed on the concrete backend client. See
+/// [`crate::routers::grpc::multimodal::capability::ensure_backend_supports_modalities`].
+fn ensure_client_supports_intermediate(
+    client: &GrpcClient,
+    intermediate: &MultimodalIntermediate,
+) -> Result<()> {
+    ensure_backend_supports_modalities(client.runtime_type(), intermediate)
 }
 
 fn assemble_sglang(

model_gateway/src/routers/grpc/multimodal/capability.rs

@@ -0,0 +1,170 @@
+//! Backend x modality capability: the single source of truth for which gRPC
+//! engine supports which input modality.
+//!
+//! Previously this was implicit and duplicated across assembly (per-backend
+//! `into_single_image_batch` / `into_single_vision_batch` / ad-hoc `bail!`s with
+//! divergent messages). Centralizing it here lets the pipeline reject an
+//! unsupported (engine, modality) request early -- at worker selection, before
+//! any media is fetched or preprocessed -- with one consistent message, and lets
+//! assembly assert against the same matrix as defense in depth.
+
+use anyhow::Result;
+use llm_multimodal::Modality;
+
+use super::MultimodalIntermediate;
+use crate::worker::RuntimeType;
+
+/// Whether `runtime` accepts multimodal inputs of `modality`.
+///
+/// This is the authoritative capability matrix. Truth, transcribed from the
+/// per-backend assembly arms:
+/// - Image: SGLang, vLLM, TRT-LLM, TokenSpeed
+/// - ImageEmbeds: TokenSpeed only (pre-computed embeddings are an EPD feature)
+/// - Video: vLLM, TokenSpeed
+/// - Audio: TokenSpeed
+/// - MLX: none
+///
+/// `Modality::ImageEmbeds` cannot actually reach the early check today because a
+/// [`crate::routers::grpc::multimodal::MediaBatch`] only ever yields
+/// Image/Video/Audio; it is kept in the matrix for correctness/defense in depth.
+pub(crate) fn runtime_supports_modality(runtime: RuntimeType, modality: Modality) -> bool {
+    match modality {
+        Modality::Image => matches!(
+            runtime,
+            RuntimeType::Sglang | RuntimeType::Vllm | RuntimeType::Trtllm | RuntimeType::TokenSpeed
+        ),
+        Modality::ImageEmbeds => matches!(runtime, RuntimeType::TokenSpeed),
+        Modality::Video => matches!(runtime, RuntimeType::Vllm | RuntimeType::TokenSpeed),
+        Modality::Audio => matches!(runtime, RuntimeType::TokenSpeed),
+    }
+}
+
+/// Reject early if the selected backend does not support every modality present
+/// in the request. Runs at worker selection, once the runtime is known but
+/// before media is fetched/preprocessed, so an unsupported combination fails
+/// fast with one clear message instead of dying deep in assembly.
+pub(crate) fn ensure_backend_supports_modalities(
+    runtime: RuntimeType,
+    intermediate: &MultimodalIntermediate,
+) -> Result<()> {
+    for batch in intermediate.batches() {
+        let modality = batch.media.modality();
+        anyhow::ensure!(
+            runtime_supports_modality(runtime, modality),
+            "backend {runtime} does not support {modality} inputs"
+        );
+    }
+    Ok(())
+}
+
+#[cfg(test)]
+mod tests {
+    use super::*;
+
+    /// The full backend x modality matrix, mirroring the pre-refactor assembly
+    /// dispatch so a behavior change surfaces here.
+    #[test]
+    fn capability_matrix_matches_backend_support() {
+        use Modality::{Audio, Image, ImageEmbeds, Video};
+        use RuntimeType::{Mlx, Sglang, TokenSpeed, Trtllm, Vllm};
+
+        // (runtime, modality, expected_supported)
+        let cases = [
+            (Sglang, Image, true),
+            (Sglang, Video, false),
+            (Sglang, Audio, false),
+            (Vllm, Image, true),
+            (Vllm, Video, true),
+            (Vllm, Audio, false),
+            (Trtllm, Image, true),
+            (Trtllm, Video, false),
+            (Trtllm, Audio, false),
+            (TokenSpeed, Image, true),
+            (TokenSpeed, Video, true),
+            (TokenSpeed, Audio, true),
+            (Mlx, Image, false),
+            (Mlx, Video, false),
+            (Mlx, Audio, false),
+            // ImageEmbeds is a TokenSpeed-only EPD feature.
+            (Sglang, ImageEmbeds, false),
+            (Vllm, ImageEmbeds, false),
+            (Trtllm, ImageEmbeds, false),
+            (TokenSpeed, ImageEmbeds, true),
+            (Mlx, ImageEmbeds, false),
+        ];
+
+        for (runtime, modality, expected) in cases {
+            assert_eq!(
+                runtime_supports_modality(runtime, modality),
+                expected,
+                "runtime={runtime} modality={modality} expected supported={expected}"
+            );
+        }
+    }
+
+    /// Non-gRPC runtimes are never routed to the multimodal gRPC path; they
+    /// support nothing here.
+    #[test]
+    fn non_grpc_runtimes_support_no_modality() {
+        for runtime in [RuntimeType::Unspecified, RuntimeType::External] {
+            for modality in [Modality::Image, Modality::Video, Modality::Audio] {
+                assert!(!runtime_supports_modality(runtime, modality));
+            }
+        }
+    }
+
+    fn single_image_intermediate() -> MultimodalIntermediate {
+        use std::{collections::HashMap, sync::Arc};
+
+        use llm_multimodal::{
+            EncoderFieldLayouts, ImageDetail, ImageFrame, ImageSource, PlaceholderRange,
+            PreprocessedEncoderInputs,
+        };
+        use ndarray::{ArrayD, IxDyn};
+
+        use super::super::{MediaBatch, PrecomputedMultimodalIntermediate, PromptBinding};
+
+        MultimodalIntermediate::try_new(vec![PrecomputedMultimodalIntermediate {
+            preprocessed: PreprocessedEncoderInputs {
+                encoder_input: ArrayD::from_shape_vec(IxDyn(&[1, 1]), vec![1.0]).unwrap(),
+                feature_token_counts: vec![1],
+                item_sizes: vec![(1, 1)],
+                model_specific: HashMap::new(),
+            },
+            media: MediaBatch::Images(vec![Arc::new(ImageFrame::new(
+                image::DynamicImage::new_rgb8(1, 1),
+                bytes::Bytes::from_static(b"image"),
+                ImageDetail::Auto,
+                ImageSource::InlineBytes,
+                "image-hash".to_string(),
+            ))]),
+            bindings: vec![PromptBinding {
+                item_index: 0,
+                prompt_ordinal: 0,
+                structural: PlaceholderRange {
+                    offset: 0,
+                    length: 1,
+                },
+                patches: vec![],
+            }],
+            placeholder_token_id: Some(10),
+            field_layouts: EncoderFieldLayouts::default(),
+            keep_on_cpu_keys: vec![],
+        }])
+        .unwrap()
+    }
+
+    /// The early check accepts a supported (backend, modality) pair and rejects
+    /// an unsupported one with the single consistent message.
+    #[test]
+    fn early_check_gates_backend_modality() {
+        let intermediate = single_image_intermediate();
+
+        // Image on SGLang is supported.
+        assert!(ensure_backend_supports_modalities(RuntimeType::Sglang, &intermediate).is_ok());
+
+        // Image on MLX is not.
+        let err = ensure_backend_supports_modalities(RuntimeType::Mlx, &intermediate).unwrap_err();
+        assert_eq!(err.to_string(), "backend mlx does not support image inputs");
+    }
+}

model_gateway/src/routers/grpc/multimodal/mod.rs

@@ -24,6 +24,7 @@ use llm_multimodal::{
 };
 
 mod assemble;
+mod capability;
 mod config;
 mod detect;
 mod pixel_cache;
@@ -36,6 +37,7 @@ pub(crate) use assemble::{
     assemble_multimodal_data, assemble_multimodal_data_after_encode,
     assemble_tokenspeed_for_encode, encode_routing_hashes,
 };
+pub(crate) use capability::ensure_backend_supports_modalities;
 pub(crate) use config::{
     load_preprocessor_config_file, load_video_preprocessor_config, MultimodalComponents,
     MultimodalConfigRegistry, MultimodalModelConfig,