[TRTLLM-11951][feat] Chunked prefill for non-contiguous multimodal data + preproc fixes

tensorrt_llm/_torch/auto_deploy/shim/ad_executor.py

@@ -37,7 +37,7 @@
 from tensorrt_llm._torch.speculative import get_spec_drafter
 from tensorrt_llm._torch.speculative.eagle3 import Eagle3OneModelSampler, Eagle3ResourceManager
 from tensorrt_llm._utils import nvtx_range
-from tensorrt_llm.inputs.multimodal import MultimodalRuntimeData
+from tensorrt_llm.inputs.multimodal import MultimodalRuntimeData, check_mm_embed_cumsum_if_needed
 from tensorrt_llm.llmapi.llm_args import (
     ContextChunkingPolicy,
     EagleDecodingConfig,
@@ -76,9 +76,16 @@
 from ..utils.logger import ad_logger
 from .interface import CachedSequenceInterface, GetInferenceModel
 
-# `layout_metadata` is a reserved multimodal payload used to carry request-level
-# layout semantics that the wrapper reconstructs into tensor kwargs separately.
-_RESERVED_MM_DATA_KEYS = frozenset({"layout_metadata"})
+# Non-tensor multimodal metadata consumed by _store_prefill_multimodal_metadata.
+# These keys must NOT leak into the generic extra_args dict — entries there
+# are expected to be tensors, and these are lists or nested dicts.
+_RESERVED_MM_DATA_KEYS = frozenset(
+    {
+        "layout_metadata",
+        "special_token_offsets",
+        "multimodal_embed_mask_cumsum",
+    }
+)
 
 
 @dataclass
@@ -406,7 +413,7 @@ def _call_func():
 class ADEngine(ModelEngine):
     """The AutoDeploy Engine (ADEngine) is the main engine interface to execute AutoDeploy models.
 
-    It follows the ``ModelEngine`` abstractions and is responsible for building the ad-optimized
+    It follows the `ModelEngine` abstractions and is responsible for building the ad-optimized
     model, converting TRT-LLM scheduled requests into ad-native (pytorch-native) inputs, running
     the model, and returning correctly formatted logits.
     """
@@ -609,9 +616,6 @@ def _store_prefill_multimodal_metadata(
                     "layout_metadata", req.py_multimodal_data
                 )
             mm_item_types = layout_metadata.get("item_types", []) if layout_metadata else []
-            special_offsets = (
-                layout_metadata.get("special_token_offsets", []) if layout_metadata else []
-            )
             mm_pos_list = list(mm_pos) if mm_pos is not None else []
             mm_len_list = list(mm_len) if mm_len is not None else []
             mm_item_types_list = list(mm_item_types)
@@ -630,32 +634,50 @@ def _store_prefill_multimodal_metadata(
             mm_item_types_flat.extend(mm_item_types_list)
             mm_token_positions_flat.extend(mm_pos_list)
             mm_token_lengths_flat.extend(mm_len_list)
+
+            mm_data = req.py_multimodal_data or {}
+            flat_cumsum = mm_data.get("multimodal_embed_mask_cumsum")
+            # special_token_offsets indices into the dense MM-token-list (which includes both embeds and specials).
+            # It does not index into the prompt-position-indexed cumsum.
+            special_offsets = list(
+                mm_data.get("special_token_offsets")
+                or (layout_metadata or {}).get("special_token_offsets", [])
+            )
             mm_special_offsets_cu_seqlen.append(
                 mm_special_offsets_cu_seqlen[-1] + len(special_offsets)
             )
-            mm_special_offsets_flat.extend(list(special_offsets))
+            mm_special_offsets_flat.extend(special_offsets)
+
             if not mm_pos or not mm_len:
                 flat_start_list.append(0)
                 count_list.append(0)
                 continue
 
+            all_prompt_tokens = req.get_tokens(0)
+            check_mm_embed_cumsum_if_needed(
+                req.py_multimodal_data,
+                begin_compute=begin_compute,
+                end_compute=end_compute,
+                prompt_len=len(all_prompt_tokens),
+            )
+            if flat_cumsum is None:
+                # Leave flat_start / count at 0 so concatenated cursors don't advance.
+                flat_start_list.append(0)
+                count_list.append(0)
+                continue
+
+            assert flat_cumsum.numel() == len(all_prompt_tokens), (
+                f"embed_mask_cumsum length {flat_cumsum.numel()} != prompt length "
+                f"{len(all_prompt_tokens)} for request {i}"
+            )
             runtime = MultimodalRuntimeData(
                 past_seen_token_num=begin_compute,
-                mm_token_positions=list(mm_pos),
-                mm_token_lengths=list(mm_len),
                 chunk_end_pos=end_compute,
-                special_token_offsets=list(special_offsets),
+                embed_mask_cumsum=flat_cumsum,
             )
-            num_unseen = runtime.num_unseen_mm_tokens
-            num_in_chunk = runtime.num_mm_tokens_in_chunk
-            num_unseen_special = runtime.num_unseen_special_tokens
-            num_special_in_chunk = runtime.num_special_tokens_in_chunk
-            total_mm_i = runtime.total_mm_tokens_in_request
-            total_special_i = runtime.total_special_tokens_in_request
-            flat_start_i = cumsum_total_mm + num_unseen - num_unseen_special
-            flat_start_list.append(flat_start_i)
-            count_list.append(num_in_chunk - num_special_in_chunk)
-            cumsum_total_mm += total_mm_i - total_special_i
+            flat_start_list.append(cumsum_total_mm + runtime.num_cached_mm_tokens)
+            count_list.append(runtime.num_mm_tokens_in_chunk)
+            cumsum_total_mm += runtime.total_embeds_in_request
 
         extra_args["mm_item_cu_seqlen"] = [
             torch.tensor(mm_item_cu_seqlen, dtype=torch.int32, device="cpu")

tensorrt_llm/_torch/models/modeling_multimodal_utils.py

@@ -1,4 +1,4 @@
-# Copyright 2024 NVIDIA CORPORATION & AFFILIATES
+# Copyright 2024-2026 NVIDIA CORPORATION & AFFILIATES
 #
 # Licensed under the Apache License, Version 2.0 (the "License");
 # you may not use this file except in compliance with the License.
@@ -55,7 +55,7 @@ def _get_uncached_multimodal_params(
                 and "multimodal_embedding" in param.multimodal_data
                 and param.multimodal_data["multimodal_embedding"] is not None):
             logger.debug(
-                f"Skipping encoder forward for param with cached multimodal_embedding"
+                "Skipping encoder forward for param with cached multimodal_embedding"
             )
             continue
 
@@ -84,8 +84,7 @@ def _cache_multimodal_embeddings(
     for param in multimodal_params:
         if param.multimodal_runtime is not None:
             embed_lengths.append(
-                param.multimodal_runtime.total_mm_tokens_in_request -
-                param.multimodal_runtime.total_special_tokens_in_request)
+                param.multimodal_runtime.total_embeds_in_request)
 
     # Validate total length matches
     total_expected = sum(embed_lengths)
@@ -183,7 +182,7 @@ def get_multimodal_embeddings(
         if (not hasattr(uncached_multimodal_params[0], 'multimodal_runtime')
                 or uncached_multimodal_params[0].multimodal_runtime is None
                 or uncached_multimodal_params[0].multimodal_runtime.
-                total_mm_tokens_in_request is None):
+                total_embeds_in_request is None):
             logger.warning(
                 "Multimodal runtime data missing or incomplete, will not cache embeddings."
             )
@@ -238,6 +237,8 @@ def find_input_mm_embeds(
         - Supports both individual batching (len(mm_embeds) == len(multimodal_params))
           and pre-concatenated batching (len(mm_embeds) == 1)
         - Handles chunked prefill by considering chunk boundaries and current chunk tokens
+        - Example: if a request has 8 MM embed rows, 2 cached rows, and 3 rows
+          in the current chunk, this keeps rows [2:5].
     """
     # Current support two batching modes:
     # 1. Pre-concatenated mm_embeds for each batch, i.e., len(mm_embeds) == 1
@@ -251,15 +252,11 @@ def find_input_mm_embeds(
         # No slicing, return the full mm_embeds
         return mm_embeds
 
-    # Calculate total tokens that need processing (both cached and current chunk)
-    total_mm_tokens = sum([
-        param.multimodal_runtime.num_mm_tokens_in_chunk -
-        param.multimodal_runtime.num_special_tokens_in_chunk
-        for param in multimodal_params if param.multimodal_runtime is not None
-    ])
+    total_mm_tokens = sum(param.multimodal_runtime.num_mm_tokens_in_chunk
+                          for param in multimodal_params
+                          if param.multimodal_runtime is not None)
 
     if total_mm_tokens == 0:
-        # No tokens need processing, return empty list
         logger.debug(
             "All multimodal tokens are cached or beyond current chunk, skipping vision encoder forward"
         )
@@ -274,26 +271,17 @@ def find_input_mm_embeds(
         runtime = param.multimodal_runtime
         if runtime is None:
             continue
-        local_start_pos = runtime.num_unseen_mm_tokens - runtime.num_unseen_special_tokens
-        local_end_pos = local_start_pos + runtime.num_mm_tokens_in_chunk - runtime.num_special_tokens_in_chunk
+        local_start_pos = runtime.num_cached_mm_tokens
+        local_end_pos = local_start_pos + runtime.num_mm_tokens_in_chunk
         slices.append(
             (current_pos + local_start_pos, current_pos + local_end_pos))
-        if len(mm_embeds
-               ) == 1:  # pre-concatenated mm_embeds, need global offset
-            current_pos += runtime.total_mm_tokens_in_request
-            current_pos -= runtime.total_special_tokens_in_request
+        if len(mm_embeds) == 1:  # pre-concatenated; advance global cursor
+            current_pos += runtime.total_embeds_in_request
 
-    sliced_mm_embeds = []
     if len(mm_embeds) == 1:
-        sliced_mm_embeds = [mm_embeds[0][start:end] for start, end in slices]
-    else:  # slice each mm_embeds individually
-        for i, (start, end) in enumerate(slices):
-            sliced_mm_embeds.append(mm_embeds[i][start:end])
-
-    if len(mm_embeds) == 1:
-        sliced_mm_embeds = [torch.cat(sliced_mm_embeds, dim=0)]
-
-    return sliced_mm_embeds
+        sliced = [mm_embeds[0][start:end] for start, end in slices]
+        return [torch.cat(sliced, dim=0)]
+    return [mm_embeds[i][start:end] for i, (start, end) in enumerate(slices)]
 
 
 def filter_mm_token_from_input_ids(
@@ -308,30 +296,25 @@ def filter_mm_token_from_input_ids(
         vocab_size: size of the model's vocabulary
         mm_token_ids: possible token ids for multimodal tokens, if known. If not known and set to None, it is assumed that the multimodal tokens are out-of-vocabulary tokens i.e. the `input_ids` contains tokens >= vocab_size that represent the multimodal tokens.
     Note:
+        Example: input_ids=[1, 55, 2, 101], vocab_size=100, and
+        mm_token_ids=[55] returns mm_token_indices=[1]; token 101 is text
+        because explicit mm_token_ids overrides the OOV fallback.
         This function involves host-device synchronization due to torch.where() (= torch.nonzero) requiring
         host allocation. The output indices reside on the same device as input_ids.
     Returns:
         text_token_indices: indices of text tokens in the input_ids
         mm_token_indices: indices of multimodal tokens in the input_ids
     """
     if mm_token_ids is None:
-        # NOTE:
-        # If mm_token_ids is None, it is assumed that the multimodal
-        # tokens are out-of-vocab tokens i.e. the `input_ids` contains
-        # tokens >= vocab_size that represent the multimodal tokens.
-        # Since mm_token_ids can be unbounded in this case,
-        # using torch.isin() may not be performant.
-        # This provides a more performant alternative while keeping
-        # the flexibility of still specifying all possible mm_token_ids,
-        # if the user wants to.
+        # If mm_token_ids is None, assume the multimodal tokens are out-of-vocab
+        # (input_ids >= vocab_size). Avoids torch.isin() over a potentially
+        # unbounded mm_token_ids set.
         mm_token_mask = input_ids >= vocab_size
-        text_token_mask = input_ids < vocab_size
     else:
         mm_token_ids = mm_token_ids.to(input_ids.device, dtype=input_ids.dtype)
         mm_token_mask = torch.isin(input_ids, mm_token_ids)
-        text_token_mask = ~mm_token_mask
     # NOTE: torch.where() enforces a host sync
-    text_token_indices = torch.where(text_token_mask)[0]
+    text_token_indices = torch.where(~mm_token_mask)[0]
     mm_token_indices = torch.where(mm_token_mask)[0]
     return text_token_indices, mm_token_indices
 
@@ -363,6 +346,9 @@ def fuse_input_embeds(
         - If (4) multimodal run, mixed batch of context and generation requests, each context request has a multimodal feature --> return only the fused input_embeds of shape [total length, hidden_dim]. For text tokens, LLM embedding layer has already run.
     Note:
         - Precedence: If kwargs provide indices (text_token_indices and mm_token_indices), those are used. If any one of them is not provided, fallback to filtering method. Sentinel-/OOV-based filtering (e.g., tokens >= vocab_size) is used only when neither index tensor and mm_token_ids is provided.
+        - Example: len(torch.cat(mm_embeds)) must match len(mm_token_indices);
+          for chunked prefill, pass only the current chunk's mm_embeds or
+          explicit indices for the active MM token positions.
         - This function may involve host-device synchronization if indices are not provided and filtering is performed. See filter_mm_token_from_input_ids for details.
     """
     if len(mm_embeds) == 0:

tensorrt_llm/_torch/models/modeling_nemotron_nano.py

@@ -998,8 +998,11 @@ def get_num_tokens_per_image(
         image: Union[Image.Image, torch.Tensor],
         **kwargs,
     ):
-        # Dynamic resolution path.
-        if self.dynamic_tiler is not None:
+        # Dynamic resolution path — only when max_num_tiles is not
+        # explicitly overridden (e.g. for video frames which use
+        # VIDEO_MAX_NUM_TILES and rely on the InternVL tiling logic
+        # in the HF processor, not the dynamic tiler).
+        if self.dynamic_tiler is not None and "max_num_tiles" not in kwargs:
             budget = self.dynamic_tiler._max_num_patches
             params, _ = self.dynamic_tiler.process_media(image, budget)
             num_image_tokens = params.num_embeddings

tensorrt_llm/_torch/models/modeling_qwen3vl.py

@@ -4,6 +4,7 @@
 
 import torch
 import torch.nn as nn
+from PIL import Image
 from transformers import AutoProcessor, AutoTokenizer, PretrainedConfig, PreTrainedModel
 from transformers.activations import ACT2FN as HF_ACT2FN
 from transformers.models.qwen3_vl.modeling_qwen3_vl import (
@@ -256,6 +257,39 @@ def get_rope_index(
 
             return position_ids, mrope_position_deltas
 
+    def get_num_tokens_per_video(
+        self,
+        *,
+        video: List[Image.Image],
+        video_grid_thw: Optional[torch.Tensor] = None,
+        **kwargs,
+    ) -> int:
+        merge = self.config.vision_config.spatial_merge_size
+        if video_grid_thw is not None:
+            t, h, w = (int(x) for x in video_grid_thw)
+            return t * (h // merge) * (w // merge)
+
+        # Must run the full processor: HF's Qwen3VLProcessor._get_num_multimodal_tokens
+        # (what the base class default delegates to) raises on video-only calls
+        # and returns a wrong-formula fallback that would break chunked prefill.
+        do_rescale = not (video and isinstance(video[0], torch.Tensor))
+        processed = self._processor(
+            text=["<|vision_start|><|video_pad|><|vision_end|>"],
+            videos=[video],
+            padding=True,
+            do_rescale=do_rescale,
+            return_tensors="pt",
+            **kwargs,
+        )
+        vgt = processed.get("video_grid_thw")
+        if vgt is None or len(vgt) == 0:
+            raise RuntimeError(
+                "get_num_tokens_per_video: HF processor returned no "
+                "video_grid_thw for the provided video."
+            )
+        t, h, w = (int(x) for x in vgt[0].tolist())
+        return t * (h // merge) * (w // merge)
+
     def _preprocess(
         self, text: Dict[str, Any], mm_data: Dict[str, Any], mm_processor_kwargs: Dict[str, Any]
     ):

tensorrt_llm/_torch/pyexecutor/llm_request.py

@@ -384,13 +384,17 @@ def append_log_probs(self,
 
     def append_mm_embeddings(self, mm_embeddings: torch.Tensor,
                              multimodal_lengths: List[int]):
-        """Split concatenated embeddings by multimodal_lengths and create handles for each.
+        """Split concatenated embeddings by per-item lengths and create handles.
 
         Args:
-            mm_embeddings: Concatenated multimodal embeddings tensor of shape [total_tokens, hidden_dim]
-            multimodal_lengths: List of token lengths for each multimodal item
+            mm_embeddings: Concatenated multimodal embeddings tensor of shape
+                [total_tokens, hidden_dim].
+            multimodal_lengths: Current per-item split lengths.
         """
-        # Split the concatenated tensor by lengths to get per-item embeddings
+        # TODO(TRTLLM-12175): callers currently pass request.multimodal_lengths,
+        # a prompt-side MM-token count that may include non-embedding
+        # special/framing tokens. This split needs per-item encoder-output
+        # embedding lengths instead.
         split_embeddings = torch.split(mm_embeddings, multimodal_lengths, dim=0)
 
         # Create a SharedTensorContainer handle for each split