feat: vLLM backend with frontend media decoding (#5781)

Signed-off-by: Alexandre Milesi <milesial@users.noreply.github.com>

Cargo.lock

@@ -2354,6 +2354,7 @@ dependencies = [
  "erased-serde",
  "etcd-client",
  "ffmpeg-next",
+ "flate2",
  "futures",
  "futures-util",
  "galil-seiferas",

components/src/dynamo/common/utils/media_nixl.py

+# SPDX-FileCopyrightText: Copyright (c) 2025-2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
+# SPDX-License-Identifier: Apache-2.0
+
+import logging
+import time
+import uuid
+from typing import Any, Dict, Tuple
+
+import numpy as np
+import torch
+
+from dynamo import nixl_connect
+from dynamo.nixl_connect import OperationKind, RdmaMetadata, SerializedDescriptor
+
+logger = logging.getLogger(__name__)
+
+
+async def read_decoded_media_via_nixl(
+    connector: nixl_connect.Connector,
+    decoded_meta: Dict[str, Any],
+    return_metadata: bool = False,
+) -> np.ndarray | Tuple[np.ndarray, Dict[str, Any]]:
+    """
+    Read pre-decoded media data via NIXL RDMA transfer, into a CPU numpy array.
+
+    Args:
+        connector: Initialized NIXL connector for RDMA operations.
+        decoded_meta: Metadata dict from the frontent, containing nixl_metadata, shape, dtype, nixl_descriptor, and metadata.
+
+    Returns:
+        np.ndarray containing the transferred media data.
+        Dict[str, Any] containing the media metadata.
+    """
+    rdma_metadata = decoded_meta["nixl_metadata"]
+    descriptor = decoded_meta["nixl_descriptor"]
+    remote_device = (
+        "cpu"
+        if descriptor.get("mem_type", "dram").lower() == "dram"
+        else f"cuda:{descriptor.get('device_id', 0)}"
+    )
+
+    rdma_metadata = RdmaMetadata(
+        descriptors=[
+            SerializedDescriptor(
+                device=remote_device,
+                ptr=descriptor["addr"],
+                size=descriptor["size"],
+            )
+        ],
+        nixl_metadata=rdma_metadata,
+        notification_key=str(uuid.uuid4()),
+        operation_kind=int(OperationKind.READ),
+    )
+
+    # Create empty tensor to receive RDMA data
+    shape = decoded_meta["shape"]
+    dtype_str = decoded_meta.get("dtype", "uint8").lower()
+    alloc_start = time.perf_counter()
+    tensor = torch.empty(shape, dtype=getattr(torch, dtype_str))
+    alloc_end = time.perf_counter()
+    local_descriptor = nixl_connect.Descriptor(tensor)
+
+    read_start = time.perf_counter()
+    read_op = await connector.begin_read(rdma_metadata, local_descriptor)
+    await read_op.wait_for_completion()
+    read_end = time.perf_counter()
+
+    logger.debug(
+        f"Loaded media via NIXL RDMA: shape={shape}, "
+        f"read_time={read_end - read_start:.4f}s, "
+        f"alloc_time={alloc_end - alloc_start:.6f}s"
+    )
+
+    array = tensor.numpy()  # zero-copy
+    if return_metadata:
+        return array, decoded_meta.get("metadata")
+    else:
+        return array

components/src/dynamo/vllm/args.py

@@ -71,6 +71,7 @@ class Config:
     enable_multimodal: bool = False
     multimodal_encode_prefill_worker: bool = False
     mm_prompt_template: str = "USER: <image>\n<prompt> ASSISTANT:"
+    frontend_decoding: bool = False
 
     # vLLM-native encoder worker (ECConnector mode)
     vllm_native_encoder_worker: bool = False
@@ -217,6 +218,15 @@ def parse_args() -> Config:
             "'USER: <image> please describe the image ASSISTANT:'."
         ),
     )
+    parser.add_argument(
+        "--frontend-decoding",
+        action="store_true",
+        help=(
+            "Enable frontend decoding of multimodal images. "
+            "When enabled, images are decoded in the Rust frontend and transferred to the backend via NIXL RDMA. "
+            "Without this flag, images are decoded in the Python backend (default behavior)."
+        ),
+    )
     parser.add_argument(
         "--vllm-native-encoder-worker",
         action="store_true",
@@ -402,6 +412,7 @@ def parse_args() -> Config:
     config.multimodal_encode_prefill_worker = args.multimodal_encode_prefill_worker
     config.enable_multimodal = args.enable_multimodal
     config.mm_prompt_template = args.mm_prompt_template
+    config.frontend_decoding = args.frontend_decoding
     config.vllm_native_encoder_worker = args.vllm_native_encoder_worker
     config.ec_connector_backend = args.ec_connector_backend
     config.ec_storage_path = args.ec_storage_path

components/src/dynamo/vllm/handlers.py

@@ -21,7 +21,9 @@ from vllm.outputs import RequestOutput
 from vllm.sampling_params import SamplingParams, StructuredOutputsParams
 from vllm.v1.engine.exceptions import EngineDeadError
 
+import dynamo.nixl_connect as nixl_connect
 from dynamo.common.utils.input_params import InputParamManager
+from dynamo.common.utils.media_nixl import read_decoded_media_via_nixl
 from dynamo.common.utils.otel_tracing import build_trace_headers
 from dynamo.llm import (
     ModelInput,
@@ -244,6 +246,7 @@ class BaseWorkerHandler(ABC):
         config=None,
         use_vllm_tokenizer: bool = False,
         shutdown_event: asyncio.Event | None = None,
+        enable_frontend_decoding: bool = False,
     ):
         self.runtime = runtime
         self.component = component
@@ -257,6 +260,10 @@ class BaseWorkerHandler(ABC):
         self.temp_dirs: list[tempfile.TemporaryDirectory] = []
         self.model_max_len = model_max_len
         self.enable_multimodal = enable_multimodal
+        self.enable_frontend_decoding = enable_frontend_decoding
+        # NIXL connector for frontend decoding - lazy initialized
+        self._nixl_connector = None
+        self._nixl_connector_lock = asyncio.Lock()
         # LoRA tracking
         self.lora_id_for_name: dict[str, int] = {}
         self.lora_name_to_path: dict[str, str] = {}
@@ -879,6 +886,10 @@ class BaseWorkerHandler(ABC):
         """
         Load a batch of images from multimodal data items.
 
+        Supports two paths:
+        1. Url variant: Download and decode image from URL (default)
+        2. Decoded variant: Read pre-decoded image via NIXL RDMA (requires --frontend-decoding)
+
         Args:
             image_mm_items: List of multimodal data items for images
         Returns:
@@ -887,25 +898,41 @@ class BaseWorkerHandler(ABC):
             Exception: If any image fails to load
         """
         image_futures = []
+
         for item in image_mm_items:
             if isinstance(item, dict) and URL_VARIANT_KEY in item:
+                # URL path: download and decode in Python backend
                 url = item[URL_VARIANT_KEY]
                 image_futures.append(self.image_loader.load_image(url))
                 logger.debug(f"Preparing to load image from URL: {url[:80]}...")
             elif isinstance(item, dict) and DECODED_VARIANT_KEY in item:
-                logger.warning(
-                    "Decoded multimodal data not yet supported in standard worker"
-                )
+                if self.enable_frontend_decoding:
+                    async with self._nixl_connector_lock:
+                        if self._nixl_connector is None:
+                            self._nixl_connector = nixl_connect.Connector()
+                            await self._nixl_connector.initialize()
+
+                    metadata = item[DECODED_VARIANT_KEY]
+                    image_futures.append(
+                        read_decoded_media_via_nixl(self._nixl_connector, metadata)
+                    )
+                else:
+                    logger.error(
+                        "Received Decoded multimodal data but --frontend-decoding not enabled. "
+                        "Use --frontend-decoding flag to enable NIXL RDMA image transfer."
+                    )
+                    raise ValueError("Could not load decoded media from frontend")
 
+        # Process images in parallel
         results = await asyncio.gather(*image_futures, return_exceptions=True)
         loaded_images = []
         collective_exceptions = ""
-        for i, result in enumerate(results):
+        for media_item, result in zip(image_mm_items, results):
             if isinstance(result, Exception):
-                url = image_mm_items[i].get(URL_VARIANT_KEY, "unknown")
-                logger.error(f"Failed to load image from {url[:80]}...: {result}")
+                source = media_item.get(URL_VARIANT_KEY, "decoded")
+                logger.error(f"Failed to load image from {source[:80]}...: {result}")
                 collective_exceptions += (
-                    f"Failed to load image from {url[:80]}...: {result}\n"
+                    f"Failed to load image from {source[:80]}...: {result}\n"
                 )
                 continue
             loaded_images.append(result)
@@ -1238,6 +1265,7 @@ class DecodeWorkerHandler(BaseWorkerHandler):
         config=None,
         use_vllm_tokenizer: bool = False,
         shutdown_event: asyncio.Event | None = None,
+        enable_frontend_decoding: bool = False,
     ):
         super().__init__(
             runtime,
@@ -1250,6 +1278,7 @@ class DecodeWorkerHandler(BaseWorkerHandler):
             config,
             use_vllm_tokenizer,
             shutdown_event,
+            enable_frontend_decoding,
         )
 
     async def generate(self, request, context):
@@ -1451,6 +1480,7 @@ class PrefillWorkerHandler(BaseWorkerHandler):
         config=None,
         use_vllm_tokenizer: bool = False,
         shutdown_event: asyncio.Event | None = None,
+        enable_frontend_decoding: bool = False,
     ):
         super().__init__(
             runtime,
@@ -1463,6 +1493,7 @@ class PrefillWorkerHandler(BaseWorkerHandler):
             config,
             use_vllm_tokenizer,
             shutdown_event,
+            enable_frontend_decoding,
         )
 
     async def generate(self, request, context):

components/src/dynamo/vllm/main.py

@@ -27,6 +27,17 @@ from dynamo.llm import (
     fetch_llm,
     register_llm,
 )
+
+# Optional imports for frontend decoding support
+try:
+    from dynamo.llm import MediaDecoder, MediaFetcher
+
+    MEDIA_DECODER_AVAILABLE = True
+except ImportError:
+    MediaDecoder = None
+    MediaFetcher = None
+    MEDIA_DECODER_AVAILABLE = False
+
 from dynamo.runtime import DistributedRuntime
 from dynamo.runtime.logging import configure_dynamo_logging
 from dynamo.vllm.multimodal_handlers import (
@@ -407,6 +418,23 @@ async def register_vllm_model(
     data_parallel_size = getattr(vllm_config.parallel_config, "data_parallel_size", 1)
     runtime_config.data_parallel_size = data_parallel_size
 
+    # Configure media decoder for frontend image decoding when enabled
+    # This enables frontend to decode images and transfer via NIXL RDMA
+    media_decoder = None
+    media_fetcher = None
+    if config.frontend_decoding:
+        if not MEDIA_DECODER_AVAILABLE:
+            raise RuntimeError(
+                "--frontend-decoding requires MediaDecoder support. "
+                "Ensure dynamo.llm module includes MediaDecoder and MediaFetcher."
+            )
+        media_decoder = MediaDecoder()
+        media_decoder.enable_image({"limits": {"max_alloc": 128 * 1024 * 1024}})
+        # media_decoder.enable_video({})
+
+        media_fetcher = MediaFetcher()
+        media_fetcher.timeout_ms(30000)
+
     await register_llm(
         model_input,
         model_type,
@@ -417,6 +445,8 @@ async def register_vllm_model(
         migration_limit=migration_limit,
         runtime_config=runtime_config,
         custom_template_path=config.custom_jinja_template,
+        media_decoder=media_decoder,
+        media_fetcher=media_fetcher,
     )
 
 
@@ -449,6 +479,7 @@ async def init_prefill(
         config=config,
         use_vllm_tokenizer=config.use_vllm_tokenizer,
         shutdown_event=shutdown_event,
+        enable_frontend_decoding=config.frontend_decoding,
     )
     handler.add_temp_dir(prometheus_temp_dir)
 
@@ -577,6 +608,7 @@ async def init(
         config=config,
         use_vllm_tokenizer=config.use_vllm_tokenizer,
         shutdown_event=shutdown_event,
+        enable_frontend_decoding=config.frontend_decoding,
     )
     handler.add_temp_dir(prometheus_temp_dir)
 

lib/bindings/python/Cargo.lock

@@ -1662,6 +1662,7 @@ dependencies = [
  "erased-serde",
  "etcd-client",
  "ffmpeg-next",
+ "flate2",
  "futures",
  "futures-util",
  "galil-seiferas",

lib/bindings/python/src/dynamo/nixl_connect/__init__.py

@@ -441,10 +441,13 @@ class ActiveOperation(AbstractOperation):
         self._status = OperationStatus.CANCELLED
         self._xfer_hndl = None
 
-    async def _wait_for_completion_(self) -> None:
+    async def _wait_for_completion_(
+        self, min_poll_ms=5, max_poll_ms=100, backoff_factor=1.5
+    ) -> None:
         # Loop until the operation is no longer in progress (or "initialized"),
         # yielding control to the event loop to allow other operations to run.
         iteration_count = 0
+        sleep_time = min_poll_ms
         while True:
             if iteration_count & 10 == 0:
                 logger.debug(
@@ -452,10 +455,9 @@ class ActiveOperation(AbstractOperation):
                 )
             match self.status:
                 # "in progress" or "initialized" means the operation is ongoing.
-                case OperationStatus.INITIALIZED:
-                    await asyncio.sleep(0.1)
-                case OperationStatus.IN_PROGRESS:
-                    await asyncio.sleep(0.1)
+                case OperationStatus.INITIALIZED | OperationStatus.IN_PROGRESS:
+                    await asyncio.sleep(sleep_time / 1000)
+                    sleep_time = min(sleep_time * backoff_factor, max_poll_ms)
                 # Any other state indicates completion or error.
                 case _:
                     return
@@ -1371,16 +1373,18 @@ class PassiveOperation(AbstractOperation):
             f")"
         )
 
-    async def _wait_for_completion_(self) -> None:
+    async def _wait_for_completion_(
+        self, min_poll_ms=5, max_poll_ms=100, backoff_factor=1.5
+    ) -> None:
         # Loop until the operation is no longer in progress (or "initialized"),
         # yielding control to the event loop to allow other operations to run.
+        sleep_time = min_poll_ms
         while True:
             match self.status:
                 # "in progress" or "initialized" means the operation is ongoing.
-                case OperationStatus.INITIALIZED:
-                    await asyncio.sleep(0.1)
-                case OperationStatus.IN_PROGRESS:
-                    await asyncio.sleep(0.1)
+                case OperationStatus.INITIALIZED | OperationStatus.IN_PROGRESS:
+                    await asyncio.sleep(sleep_time / 1000)
+                    sleep_time = min(sleep_time * backoff_factor, max_poll_ms)
                 # Any other state indicates completion or error.
                 case _:
                     return

lib/llm/Cargo.toml

@@ -25,7 +25,7 @@ block-manager = ["dep:nixl-sys", "dep:cudarc", "dep:nix", "dep:aligned-vec", "de
 block-manager-bench = ["block-manager", "testing-full", "dep:clap", "dep:indicatif"]
 cuda = ["dep:cudarc"]
 integration = ["dynamo-runtime/integration"]
-media-nixl = ["dep:nixl-sys", "dep:dynamo-memory"]
+media-nixl = ["dep:nixl-sys", "dep:dynamo-memory", "dep:flate2"]
 media-ffmpeg = ["dep:video-rs", "dep:ffmpeg-next", "dep:memfile", "media-nixl"]
 kv-router-stress = ["dep:clap", "dep:indicatif"]
 
@@ -110,6 +110,9 @@ nixl-sys = { version = "=0.9.0", optional = true }
 cudarc = { workspace = true, optional = true }
 nix = { version = "0.26", optional = true }
 
+# media-nixl (zlib compression for NIXL metadata)
+flate2 = { version = "1", optional = true }
+
 # block_manager_bench
 clap = { version = "4.5.49", features = ["derive"], optional = true }
 indicatif = { version = "0.18.0", optional = true }

lib/llm/src/preprocessor/media/rdma.rs

@@ -10,6 +10,8 @@ use {
     base64::{Engine as _, engine::general_purpose},
     dynamo_memory::SystemStorage,
     dynamo_memory::nixl::{self, NixlAgent, NixlDescriptor, RegisteredView},
+    flate2::{Compression, write::ZlibEncoder},
+    std::io::Write,
     std::sync::Arc,
 };
 
@@ -108,7 +110,8 @@ impl<D: Dimension> TryFrom<ArrayBase<OwnedRepr<u8>, D>> for DecodedMediaData {
 }
 
 // Get NIXL metadata for a descriptor
-// Avoids cross-request leak possibility and reduces metadata size
+// Returns zlib-compressed, base64-encoded metadata in format: "b64:<compressed_base64>"
+// This format matches what Python nixl_connect expects for RdmaMetadata.nixl_metadata
 // TODO: pre-allocate a fixed NIXL-registered RAM pool so metadata can be cached on the target?
 #[cfg(feature = "media-nixl")]
 pub fn get_nixl_metadata(agent: &NixlAgent, _storage: &SystemStorage) -> Result<String> {
@@ -118,7 +121,12 @@ pub fn get_nixl_metadata(agent: &NixlAgent, _storage: &SystemStorage) -> Result<
     // reg_desc_list.add_storage_desc(storage)?;
     // let nixl_partial_md = agent.raw_agent().get_local_partial_md(&reg_desc_list, None)?;
 
-    let b64_encoded = general_purpose::STANDARD.encode(&nixl_md);
+    // Compress with zlib (level 6, matching Python's default)
+    let mut encoder = ZlibEncoder::new(Vec::new(), Compression::new(6));
+    encoder.write_all(&nixl_md)?;
+    let compressed = encoder.finish()?;
+
+    let b64_encoded = general_purpose::STANDARD.encode(&compressed);
     Ok(format!("b64:{}", b64_encoded))
 }
 

tests/serve/test_vllm.py

@@ -304,6 +304,37 @@ vllm_configs = {
             )
         ],
     ),
+    "multimodal_agg_frontend_decoding": VLLMConfig(
+        name="multimodal_agg_frontend_decoding",
+        directory=vllm_dir,
+        script_name="agg_multimodal.sh",
+        marks=[pytest.mark.gpu_1, pytest.mark.pre_merge],
+        model="Qwen/Qwen2-VL-2B-Instruct",
+        # Pass --frontend-decoding to enable Rust frontend image decoding + NIXL RDMA transfer
+        script_args=[
+            "--model",
+            "Qwen/Qwen2-VL-2B-Instruct",
+            "--frontend-decoding",
+        ],
+        request_payloads=[
+            chat_payload(
+                [
+                    {
+                        "type": "text",
+                        "text": "What colors are in the following image? Respond only with the colors.",
+                    },
+                    {
+                        "type": "image_url",
+                        "image_url": {"url": MULTIMODAL_IMG_URL},
+                    },
+                ],
+                repeat_count=1,
+                expected_response=["green"],
+                temperature=0.0,
+                max_tokens=100,
+            )
+        ],
+    ),
     "multimodal_agg_llava_epd": VLLMConfig(
         name="multimodal_agg_llava_epd",
         directory=vllm_dir,