[Inference] Fix an unsafe fork for frontend replicas and support Nemotron VL vision preprocessing. (#7556)

Signed-off-by: Cory Ye <cye@nvidia.com>
Co-authored-by: Philip Petrakian <ppetrakian@nvidia.com>

Author: cspades

megatron/core/inference/config.py

@@ -1,7 +1,7 @@
 # Copyright (c) 2026 NVIDIA CORPORATION & AFFILIATES. All rights reserved.
 
 import warnings
-from dataclasses import InitVar, dataclass, field
+from dataclasses import InitVar, dataclass, field, replace
 from enum import Enum
 from typing import List, Literal, Optional, Tuple
 
@@ -282,6 +282,31 @@ class ImageProcessingConfig:
     max_num_tiles: int = 1
     use_thumbnail: bool = False
     num_img_embeddings_per_tile: int = 0
+    dynamic_resolution_model_length: Optional[int] = None
+    """Model-length budget used by processors that divide capacity across images."""
+    dynamic_resolution_rounding_mode: Literal["ceil", "round_plus_half"] = "ceil"
+    """Patch-grid rounding contract: ``ceil`` or ``round_plus_half``."""
+    dynamic_resolution_resize_mode: Literal["pil", "torch_bicubic_antialias"] = "pil"
+    """Resize contract: ``pil`` or ``torch_bicubic_antialias``."""
+
+    def __post_init__(self):
+        if self.dynamic_resolution_rounding_mode not in ("ceil", "round_plus_half"):
+            raise ValueError(
+                "ImageProcessingConfig.dynamic_resolution_rounding_mode must be "
+                "'ceil' or 'round_plus_half'."
+            )
+        if self.dynamic_resolution_resize_mode not in ("pil", "torch_bicubic_antialias"):
+            raise ValueError(
+                "ImageProcessingConfig.dynamic_resolution_resize_mode must be "
+                "'pil' or 'torch_bicubic_antialias'."
+            )
+        if (
+            self.dynamic_resolution_model_length is not None
+            and self.dynamic_resolution_model_length <= 4
+        ):
+            raise ValueError(
+                "ImageProcessingConfig.dynamic_resolution_model_length must be " "greater than 4."
+            )
 
 
 @dataclass
@@ -304,6 +329,24 @@ class MediaPromptSpec:
     prefix: str = ""
     suffix: str = ""
     input_marker: Optional[str] = None
+    content_part_separator: str = ""
+    expansion_mode: Literal["single", "temporal_patch"] = "single"
+    include_frame_timestamps_for_nemotron_vl: bool = False
+
+    def __post_init__(self):
+        if self.expansion_mode not in ("single", "temporal_patch"):
+            raise ValueError(
+                "MediaPromptSpec.expansion_mode must be 'single' or "
+                f"'temporal_patch', got {self.expansion_mode!r}."
+            )
+        if (
+            self.include_frame_timestamps_for_nemotron_vl
+            and self.expansion_mode != "temporal_patch"
+        ):
+            raise ValueError(
+                "MediaPromptSpec.include_frame_timestamps_for_nemotron_vl requires "
+                "expansion_mode='temporal_patch'."
+            )
 
 
 @dataclass(frozen=True)
@@ -322,13 +365,14 @@ def get_spec(self, modality: str) -> MediaPromptSpec:
         raise ValueError(f"Unsupported media modality: {modality!r}")
 
     @classmethod
-    def from_dict(cls, value):
-        """Build from image and video specs."""
+    def from_dict(cls, value, defaults=None):
+        """Build from image and video overrides, preserving optional defaults."""
         if not value:
-            return cls()
+            return defaults or cls()
+        defaults = defaults or cls()
         return cls(
-            image_spec=MediaPromptSpec(**value.get("image_spec", {})),
-            video_spec=MediaPromptSpec(**value.get("video_spec", {})),
+            image_spec=replace(defaults.image_spec, **dict(value.get("image_spec", {}))),
+            video_spec=replace(defaults.video_spec, **dict(value.get("video_spec", {}))),
         )
 
 
@@ -478,6 +522,9 @@ class InferenceConfig:
     video_preprocessing_config: Optional[VideoProcessingConfig] = None
     """Configuration for decoding and preprocessing raw video payloads."""
 
+    multimodal_prompt_config: Optional[MultimodalPromptConfig] = None
+    """Optional per-engine overrides for the inference wrapper's media prompt contract."""
+
     use_flashinfer_fused_rope: Optional[bool] = False
     """
     If True, use flashinfer's fused rope implementation.

megatron/core/inference/engines/dynamic_engine.py

@@ -182,6 +182,8 @@ class _VisionCacheEntry:
     num_img_embeddings_per_tile: int
     imgs_sizes: Optional[Tensor]
     num_frames: Optional[Tensor]
+    video_frame_indices: Optional[List[List[int]]] = None
+    video_fps: Optional[List[float]] = None
 
 
 def format_mem_bytes(mem_bytes):
@@ -733,6 +735,14 @@ def _refresh_vlm_request_data(self, request: DynamicVLMInferenceRequest) -> None
             expansion_kwargs = {"num_tiles": num_tiles, "imgs_sizes": imgs_sizes}
             if num_frames is not None:
                 expansion_kwargs["num_frames"] = num_frames
+                prompt_config = wrapper.multimodal_prompt_config
+                if (
+                    prompt_config is not None
+                    and prompt_config.video_spec.expansion_mode == "temporal_patch"
+                ):
+                    expansion_kwargs["tokenizer"] = self.controller.tokenizer
+                    expansion_kwargs["video_frame_indices"] = request.video_frame_indices
+                    expansion_kwargs["video_fps"] = request.video_fps
             _, mask_list = wrapper.expand_image_tokens(
                 [request.compact_prompt_tokens.tolist()],
                 image_token_id=media_token_id,
@@ -776,6 +786,8 @@ def _refresh_vlm_request_data(self, request: DynamicVLMInferenceRequest) -> None
             num_img_embeddings_per_tile=request.num_img_embeddings_per_tile,
             imgs_sizes=imgs_sizes,
             num_frames=num_frames,
+            video_frame_indices=request.video_frame_indices,
+            video_fps=request.video_fps,
         )
         self.context.add_vlm_request_data(
             request.request_id, image_embeddings=embeddings, image_token_mask=mask
@@ -812,6 +824,8 @@ def _cache_vision_embedding(
         num_img_embeddings_per_tile: int = 0,
         imgs_sizes: Optional[Tensor] = None,
         num_frames: Optional[Tensor] = None,
+        video_frame_indices: Optional[List[List[int]]] = None,
+        video_fps: Optional[List[float]] = None,
     ) -> None:
         if not cache_key or self.vision_embedding_cache_max_bytes == 0:
             return
@@ -823,6 +837,8 @@ def _cache_vision_embedding(
             num_img_embeddings_per_tile=num_img_embeddings_per_tile,
             imgs_sizes=imgs_sizes,
             num_frames=num_frames,
+            video_frame_indices=video_frame_indices,
+            video_fps=video_fps,
         )
         cache_entry_bytes = self._vision_cache_entry_nbytes(entry)
         if cache_entry_bytes > self.vision_embedding_cache_max_bytes:
@@ -1570,12 +1586,27 @@ def _serialize_finished_request(
                 finished_metadata=finished_metadata,
                 offload_params=request.offload_params,
             )
-        return request.serialize(
+        serialized = request.serialize(
             payload_offloaded=stage_result is not None,
             payload_stage_metadata=(
                 stage_result.response_metadata if stage_result is not None else None
             ),
         )
+        if isinstance(request, DynamicVLMInferenceRequest):
+            # Active VLM requests retain these inputs for refresh/cache bookkeeping,
+            # but completion replies must not echo large input-only tensors.
+            for key in (
+                "imgs",
+                "num_tiles",
+                "imgs_sizes",
+                "num_frames",
+                "video_frame_indices",
+                "video_fps",
+                "image_embeddings",
+                "image_token_mask",
+            ):
+                serialized[key] = None
+        return serialized
 
     def _handle_failed_request(self, request_id: int):
         """Handle a failed request by sending the reply immediately.
@@ -1842,6 +1873,8 @@ def add_request(
         num_img_embeddings_per_tile: int = 0,
         imgs_sizes: Optional[Tensor] = None,
         num_frames: Optional[Tensor] = None,
+        video_frame_indices: Optional[List[List[int]]] = None,
+        video_fps: Optional[List[float]] = None,
         media_tokens_preexpanded: bool = False,
         offload_params: Optional[Dict] = None,
         media_cache_key: Optional[str] = None,
@@ -1875,6 +1908,9 @@ def add_request(
             imgs_sizes (Optional[Tensor]): Per-image sizes [N, 2] with [H, W].
                 Dynamic resolution.
             num_frames (Optional[Tensor]): Number of frames per image/video item.
+            video_frame_indices (Optional[List[List[int]]]): Source indices for
+                each sampled video frame.
+            video_fps (Optional[List[float]]): Source FPS for each video item.
             media_tokens_preexpanded (bool): Whether prompt token IDs already contain
                 one model token per projected media embedding.
             offload_params (Optional[Dict]): Opaque metadata forwarded to the payload stager.
@@ -1952,6 +1988,8 @@ def add_request(
                     imgs_sizes=imgs_sizes,
                     precomputed_block_hashes=precomputed_block_hashes,
                     num_frames=num_frames,
+                    video_frame_indices=video_frame_indices,
+                    video_fps=video_fps,
                     media_tokens_preexpanded=media_tokens_preexpanded,
                     offload_params=offload_params,
                     media_cache_key=media_cache_key,
@@ -2016,6 +2054,8 @@ def _build_vlm_request(
         imgs_sizes: Optional[Tensor],
         precomputed_block_hashes: Optional[List[int]] = None,
         num_frames: Optional[Tensor] = None,
+        video_frame_indices: Optional[List[List[int]]] = None,
+        video_fps: Optional[List[float]] = None,
         media_tokens_preexpanded: bool = False,
         offload_params: Optional[Dict] = None,
         media_cache_key: Optional[str] = None,
@@ -2031,6 +2071,8 @@ def _build_vlm_request(
             num_img_embeddings_per_tile = cached_vision_entry.num_img_embeddings_per_tile
             imgs_sizes = cached_vision_entry.imgs_sizes
             num_frames = cached_vision_entry.num_frames
+            video_frame_indices = cached_vision_entry.video_frame_indices
+            video_fps = cached_vision_entry.video_fps
         elif num_frames is not None:
             modality = "video"
             missing = [
@@ -2084,6 +2126,8 @@ def _build_vlm_request(
                 ("num_tiles", num_tiles),
                 ("imgs_sizes", imgs_sizes),
                 ("num_frames", num_frames),
+                ("video_frame_indices", video_frame_indices),
+                ("video_fps", video_fps),
             ):
                 if value is not None:
                     media_inputs[name] = value
@@ -2132,6 +2176,14 @@ def _build_vlm_request(
                 expansion_kwargs = {"num_tiles": num_tiles, "imgs_sizes": imgs_sizes}
                 if num_frames is not None:
                     expansion_kwargs["num_frames"] = num_frames
+                    prompt_config = inference_wrapper.multimodal_prompt_config
+                    if (
+                        prompt_config is not None
+                        and prompt_config.video_spec.expansion_mode == "temporal_patch"
+                    ):
+                        expansion_kwargs["tokenizer"] = self.controller.tokenizer
+                        expansion_kwargs["video_frame_indices"] = video_frame_indices
+                        expansion_kwargs["video_fps"] = video_fps
                 expanded_tokens_list, mask_list = inference_wrapper.expand_image_tokens(
                     token_list, image_token_id=media_token_id, **expansion_kwargs
                 )
@@ -2188,6 +2240,8 @@ def _build_vlm_request(
                     num_img_embeddings_per_tile=num_img_embeddings_per_tile,
                     imgs_sizes=imgs_sizes,
                     num_frames=num_frames,
+                    video_frame_indices=video_frame_indices,
+                    video_fps=video_fps,
                 )
 
         self.context.add_vlm_request_data(
@@ -2223,6 +2277,8 @@ def _build_vlm_request(
             num_tiles=num_tiles,
             imgs_sizes=imgs_sizes,
             num_frames=num_frames,
+            video_frame_indices=video_frame_indices,
+            video_fps=video_fps,
             media_tokens_preexpanded=media_tokens_preexpanded,
             media_cache_key=media_cache_key,
             decoder_seq_length=0,

megatron/core/inference/inference_client.py

@@ -269,6 +269,7 @@ def _pack_submit_frames(
             list: The frames to send, in wire order.
         """
         media_meta, media_payload = split_multimodal_data(
+            # If multi_modal_data is already serialized, then this is an identity function.
             serialize_multimodal_data(multi_modal_data)
         )
         return [

megatron/core/inference/inference_request.py

@@ -176,7 +176,8 @@ def serialize_multimodal_data(multi_modal_data: Any) -> Optional[Dict[str, Any]]
     Video:
         ``"video"`` accepts raw video bytes, a list of raw video bytes, or a
         preprocessed tensor dictionary containing ``imgs``, ``imgs_sizes``,
-        and ``num_frames``.
+        and ``num_frames``. Preprocessed video dictionaries may also include
+        ``video_frame_indices`` and ``video_fps`` timing metadata.
     Audio:
         Audio does not yet have any supported data preprocessing or modeling
         formats.
@@ -236,6 +237,10 @@ def serialize_multimodal_data(multi_modal_data: Any) -> Optional[Dict[str, Any]]
             wire[key] = serialize_tensor(value)
         if "num_img_embeddings_per_tile" in modality_data:
             wire["num_img_embeddings_per_tile"] = int(modality_data["num_img_embeddings_per_tile"])
+        if modality == "video":
+            for key in ("video_frame_indices", "video_fps"):
+                if key in modality_data:
+                    wire[key] = copy.deepcopy(modality_data[key])
         return {modality: wire, "media_cache_key": media_cache_key, **metadata} if wire else None
     else:
         raise TypeError(
@@ -403,6 +408,10 @@ def resolve_multimodal_data_for_engine(
             kwargs[key] = value if isinstance(value, torch.Tensor) else deserialize_tensor(value)
     if "num_img_embeddings_per_tile" in modality_data:
         kwargs["num_img_embeddings_per_tile"] = int(modality_data["num_img_embeddings_per_tile"])
+    if modality == "video":
+        for key in ("video_frame_indices", "video_fps"):
+            if key in modality_data:
+                kwargs[key] = copy.deepcopy(modality_data[key])
 
     if modality == "image":
         # Reject incomplete static-tiling payloads. Static tiling (imgs +
@@ -1451,5 +1460,7 @@ class DynamicVLMInferenceRequest(DynamicInferenceRequest, VLMInferenceRequest):
     image_token_mask: Optional[torch.Tensor] = None  # 1D, -1=text, >=0=image index
     imgs_sizes: Optional[torch.Tensor] = None
     num_frames: Optional[torch.Tensor] = None
+    video_frame_indices: Optional[List[List[int]]] = None
+    video_fps: Optional[List[float]] = None
     media_tokens_preexpanded: bool = False
     media_cache_key: Optional[str] = None

megatron/core/inference/model_inference_wrappers/abstract_model_inference_wrapper.py

@@ -68,6 +68,11 @@ def __init__(
         self.sequence_parallel = self.config.sequence_parallel
 
         self.inference_context = inference_context
+        configured_prompt_contract = getattr(
+            self.inference_context.config, "multimodal_prompt_config", None
+        )
+        if configured_prompt_contract is not None:
+            self.multimodal_prompt_config = configured_prompt_contract
 
         # Get the inference pg_collection from the config if it exists; otherwise the training
         # pg_collection might be used during RL