From 4e9f543448fb9a5dc477ca492845725ca8aedfa5 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 26 Jun 2025 23:54:04 +0900
Subject: [PATCH 01/14] feat: add MultimodalParams & putting all multimodal
 params into it and refactor HyperCLOVAX & Qwen2/2.5-VL

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_hyperclovax.py     |  97 ++++------
 .../_torch/models/modeling_qwen2vl.py         | 180 +++++++++++-------
 tensorrt_llm/_torch/pyexecutor/llm_request.py |   7 +-
 .../_torch/pyexecutor/model_engine.py         |   3 +-
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  15 +-
 tensorrt_llm/executor/executor.py             |   6 +-
 tensorrt_llm/executor/request.py              |   2 +
 tensorrt_llm/executor/worker.py               |   4 +
 tensorrt_llm/llmapi/llm.py                    |   5 +
 9 files changed, 178 insertions(+), 141 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_hyperclovax.py b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
index 7134ceeecb78..e70c363a5b64 100644
--- a/tensorrt_llm/_torch/models/modeling_hyperclovax.py
+++ b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
@@ -22,9 +22,6 @@
 from .modeling_multimodal_utils import fuse_input_embeds
 from .modeling_utils import register_auto_model
 
-DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
-
-
 # Copied from HyperCLOVAX-SEED-Vision-Instruct-3B/modeling_hyperclovax.py
 def select_best_resolution(original_size: tuple,
                            possible_resolutions: list) -> tuple:
@@ -423,11 +420,8 @@ def __init__(self,
             model_path,
             trust_remote_code=trust_remote_code,
             use_fast=self.use_fast)
-        self.tllm_image_token_id = self.pretrained_config.language_config[
+        self.tllm_multimodal_token_id = self.pretrained_config.language_config[
             "vocab_size"] + 1
-        if DISAGG:
-            self.mm_encoder = HCXVisionModel(self.pretrained_config,
-                                             skip_processor=True)
 
     def _post_process(self,
                       input_ids: torch.Tensor,
@@ -476,7 +470,7 @@ def _post_process(self,
                     batch_idx,
                     input_start + token_len:input_start + token_len +
                     vision_query_lengths[batch_idx][multi_img_idx],
-                ] = self.tllm_image_token_id
+                ] = self.tllm_multimodal_token_id
 
                 input_start += token_len + vision_query_lengths[batch_idx][
                     multi_img_idx]
@@ -531,49 +525,34 @@ def __call__(
         if not preprocessed_image:
             return fused_input_ids.to(torch.int32).tolist(), {}
 
-        if DISAGG:
-            mm_embeds = self.mm_encoder.forward(preprocessed_image)
-            mm_embeds = torch.cat(mm_embeds, dim=0)
-        else:
-            # NOTE: For now, I am using "mm_embeding" in tensor format to send the image data to the model.
-            # CASE 1: Sending raw image data
-            if isinstance(images[0], Image.Image):
-                images = [torch.from_numpy(np.array(image)) for image in images]
-            mm_embeds = torch.stack(images, dim=0)
-
-            # NOTE: After refactoring the llmRequest, we can use preprocessed_image['pixel_values'] to send the image data to the model.
-            # CASE 2: Sending preprocessed image data
-            # mm_embeds = torch.cat(preprocessed_image['pixel_values'][0],
-            #                                dim=0)
-
+        mm_data = {}
+        mm_data["image"] = {
+            "pixel_values": torch.stack(preprocessed_image['pixel_values'][0], dim=0), #TODO change the pixel_values into the Shared Tensor
+            "image_sizes": preprocessed_image.get('image_sizes', None),
+            "is_videos": preprocessed_image.get('is_videos', None),
+            "num_queries_vis_abstractors": preprocessed_image.get('num_queries_vis_abstractors', None),
+            "num_queries_vis_abstractors_slow": preprocessed_image.get('num_queries_vis_abstractors_slow', None),
+            "first_last_frames_slows": preprocessed_image.get('first_last_frames_slows', None),
+        }
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mm_embedding": mm_embeds,
+            "mm_data": mm_data
         }
 
 
 class HCXVisionModel:
 
     def __init__(self,
-                 pretrained_config: PretrainedConfig,
-                 skip_processor: bool = False):
-
+                 pretrained_config: PretrainedConfig):
+    
         self.pretrained_config = pretrained_config
         self.vision_config = self.pretrained_config.vision_config
 
         model_path = self.pretrained_config._name_or_path
 
-        # TODO: Remove this when we refactor LlmRequest
-        # NOTE: trust_remote_code can be removed once we refactor LlmRequest
-        self.skip_processor = skip_processor
-        if not self.skip_processor:
-            self.processor = AutoProcessor.from_pretrained(
-                model_path, trust_remote_code=True, use_fast=True)
-
         # NOTE: There is no way of importing mm_projector, HCXVisionCAbstractor from HF. So, can not do the sharded_loading.
         # NOTE: trust_rmemote_code can be removed once we change the model into TRT-LLM's format
         model = transformers.AutoModelForCausalLM.from_pretrained(
-            model_path, trust_remote_code=True)
-        model.eval()
+            model_path, trust_remote_code=True).eval()
         self.device = 'cuda'
 
         # TODO: Convert to TRT-LLM's SIGLIP
@@ -635,25 +614,26 @@ def _preprocess(self, mm_data: List[Any]) -> Dict[str, List[Any]]:
             for key in preprocessed_image_list[0].keys()
         }
 
-    def forward(self, mm_data: Union[List[Any], Dict[str, Any]]):
-        if not self.skip_processor:
-            # NOTE: This should be done in the input processor and got the preprocessed_image metadata from request level.
-            # But before refactoring the llmRequest, we are re-doing inputprocessor here.
-            preprocessed_image = self._preprocess(mm_data)
-        else:
-            # NOTE: When we refactor the llmRequest, we will get the extra_mm_data from mm_data, and need to make it as preprocessed_image.
-            preprocessed_image = mm_data
-            preprocessed_image["pixel_values"] = self._to_device(
-                preprocessed_image["pixel_values"])
+    def _parse_and_batch_mm_data(self, mm_data: List[Dict[str, Any]]) -> Tuple[List[torch.Tensor], Dict[str, List[Any]]]:
+        pixel_values = [list(torch.unbind(data["image"]["pixel_values"], dim=0)) for data in mm_data]
+        mm_extra_data = {
+            key: [d["image"][key][0] for d in mm_data]
+            for key in mm_data[0]["image"].keys()
+        }
+        return pixel_values, mm_extra_data
 
-        pixel_values = preprocessed_image.get("pixel_values", None)
-        image_sizes = preprocessed_image.get("image_sizes", None)
-        is_videos = preprocessed_image.get("is_videos", None)
-        num_queries_vis_abstractors = preprocessed_image.get(
+    @torch.inference_mode()
+    def forward(self, mm_data: List[Dict[str, Any]]):
+        
+        pixel_values, mm_extra_data = self._parse_and_batch_mm_data(mm_data)
+        pixel_values = self._to_device(pixel_values)
+        image_sizes = mm_extra_data.get("image_sizes", None)
+        is_videos = mm_extra_data.get("is_videos", None)
+        num_queries_vis_abstractors = mm_extra_data.get(
             "num_queries_vis_abstractors", None)
-        num_queries_vis_abstractors_slow = preprocessed_image.get(
+        num_queries_vis_abstractors_slow = mm_extra_data.get(
             "num_queries_vis_abstractors_slow", None)
-        first_last_frames_slows = preprocessed_image.get(
+        first_last_frames_slows = mm_extra_data.get(
             "first_last_frames_slows", None)
 
         len_pixel_values = [len(pixel_value) for pixel_value in pixel_values]
@@ -791,8 +771,7 @@ def __init__(self, model_config: ModelConfig):
         if hasattr(self, "llm"):
             return
 
-        if not DISAGG:
-            self.mm_encoder = HCXVisionModel(model_config.pretrained_config)
+        self.mm_encoder = HCXVisionModel(model_config.pretrained_config)
 
         llm_model_config = copy.deepcopy(model_config)
         llm_model_config.pretrained_config = PretrainedConfig.from_dict(
@@ -843,17 +822,13 @@ def forward(
             f"num_context_requests: {num_context_requests}, num_generation_requests: {num_generation_requests}"
         )
 
-        mm_data = kwargs.get("multi_modal_data", [])
+        mm_data = kwargs.get("mm_data", [])
         mm_embeds = []
         if len(mm_data) > 0:
             assert len(
                 mm_data
-            ) == num_context_requests, f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch."
-            if DISAGG:
-                # NOTE: In the DISAGG, we are assuming we get the mm_embeds from the llmRequest.
-                mm_embeds = mm_data
-            else:
-                mm_embeds = self.mm_encoder.forward(mm_data)
+            ) == num_context_requests == len(mm_data), f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch."
+            mm_embeds = self.mm_encoder.forward(mm_data)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
                                                     input_ids, mm_embeds)
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index c52b260e4cc3..a77c0697a02e 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -28,25 +28,17 @@ def __init__(self,
                  trust_remote_code: bool = True):
         self.model_config = model_config
         self.tokenizer = tokenizer
+        # TODO: change to True and also change the acoording test result
         self.use_fast = False
+        self.device = 'cuda'
         self.processor = AutoProcessor.from_pretrained(
             model_path,
             use_fast=self.use_fast,
             trust_remote_code=trust_remote_code)
 
-        # NOTE: Using attn_implementation='flash_attention_2' to avoid the issue of vision model's GPU OOM.
-        model = self.get_model_class().from_pretrained(
-            model_path,
-            torch_dtype=model_config.torch_dtype,
-            attn_implementation='flash_attention_2')
-        self.device = 'cuda'
-        self.visual = model.visual.to(self.device)
+        self.tllm_multimodal_token_id = self.model_config.vocab_size + 1
         self._post_init_()
 
-    @classmethod
-    def get_model_class(cls) -> type[PreTrainedModel]:
-        raise NotImplementedError()
-
     @classmethod
     def get_rope_index(
         cls,
@@ -284,34 +276,13 @@ def _preprocess(self, text: dict[str, any], mm_data: dict[str, any],
                               return_tensors='pt',
                               **mm_processor_kwargs)
 
-    def _process(self, pixel_values: torch.Tensor,
-                 pixel_values_videos: torch.Tensor,
-                 image_grid_thw: torch.Tensor,
-                 video_grid_thw: torch.Tensor) -> torch.Tensor:
-        embeds = []
-
-        if pixel_values is not None:
-            pixel_values = pixel_values.to(self.visual.dtype)
-            embeds.append(self.visual(pixel_values, grid_thw=image_grid_thw))
-
-        if pixel_values_videos is not None:
-            pixel_values_videos = pixel_values_videos.to(self.visual.dtype)
-            embeds.append(
-                self.visual(pixel_values_videos, grid_thw=video_grid_thw))
-
-        if embeds:
-            return torch.cat(embeds, dim=1)
-        return None
-
     def _postprocess(self, input_ids: torch.IntTensor) -> torch.IntTensor:
-        # NOTE: Qwen2-VL's input processor is doing all the work for fusing input_ids with mm_tokens. So, we just replace mm_tokens with expanded out-of-vocab ids
-
+        # NOTE: Qwen2-VL's input processor is doing all the work for fusing input_ids with mm_tokens. 
+        # So, we just replace mm_tokens with expanded out-of-vocab ids
         masks = (input_ids == self.model_config.image_token_id) | (
             input_ids == self.model_config.vision_token_id) | (
                 input_ids == self.model_config.video_token_id)
-        cumulative_counts = masks.cumsum(dim=-1)
-        values = (self.model_config.vocab_size - 1) + cumulative_counts
-        input_ids[masks] = values[masks]
+        input_ids[masks] = self.tllm_multimodal_token_id
         return input_ids
 
     def get_mrope_config(
@@ -363,17 +334,25 @@ def __call__(
 
         processed_inputs = self._preprocess(text_prompt, mm_data,
                                             mm_processor_kwargs).to(self.device)
-        if mm_data:
-            mm_features = self._process(
-                processed_inputs.get('pixel_values', None),
-                processed_inputs.get('pixel_values_videos', None),
-                processed_inputs.get('image_grid_thw', None),
-                processed_inputs.get('video_grid_thw', None))
-        else:
-            mm_features = None
+        
+        if not mm_data:
+            fused_input_ids = processed_inputs['input_ids']
+            return fused_input_ids.to(torch.int32).tolist(), {}
+
+        pixel_values = processed_inputs.get('pixel_values', None)
+        pixel_values_videos = processed_inputs.get('pixel_values_videos', None)
+        assert pixel_values is not None or pixel_values_videos is not None, "No multimodal data found"
+        
+        mm_data = {}
+        if pixel_values is not None:
+            mm_data["image"] = {"pixel_values": pixel_values, 
+                                "image_grid_thw": processed_inputs.get('image_grid_thw')}
+        if pixel_values_videos is not None:
+            mm_data["video"] = {"pixel_values_videos": pixel_values_videos, 
+                                "video_grid_thw": processed_inputs.get('video_grid_thw')}
 
         input_ids = processed_inputs['input_ids']
-
+        # TODO: We can move this to the LLM-side.
         mrope_config = self.get_mrope_config(
             input_ids, processed_inputs.get('image_grid_thw', None),
             processed_inputs.get('video_grid_thw', None),
@@ -383,23 +362,79 @@ def __call__(
         fused_input_ids = self._postprocess(input_ids[0])
 
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mm_embedding": mm_features,
-            "mrope_config": mrope_config
+            "mrope_config": mrope_config,
+            "mm_data": mm_data,
         }
 
+class Qwen2VisionModelBase:
+    def __init__(self, pretrained_config: PretrainedConfig, model_class: type[PreTrainedModel]):
+        self.pretrained_config = pretrained_config
+        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
 
-class Qwen2VLInputProcessor(Qwen2VLInputProcessorBase):
-
-    @classmethod
-    def get_model_class(cls):
-        return Qwen2VLForConditionalGeneration
+        model_path = self.pretrained_config._name_or_path
+        # TODO: Change the model class to TRT-LLM's Qwen2VisionModel
+        # NOTE: Using attn_implementation='flash_attention_2' to avoid the issue of vision model's GPU OOM.
+        model = model_class.from_pretrained(
+            model_path,
+            torch_dtype=self.pretrained_config.torch_dtype,
+            attn_implementation='flash_attention_2').eval()
+        self.visual = model.visual.to(self.device)
+    
+    def _parse_and_batch_mm_data(self, mm_data: List[Dict[str, Any]]) -> Tuple[Dict[str, Any], Dict[str, List[Any]]]:
+
+        pixel_values_list = []
+        pixel_values_videos_list = []
+        image_grid_thw_list = []
+        video_grid_thw_list = []
+        
+        for mm_content in mm_data:
+            # Process images if present
+            if "image" in mm_content and mm_content["image"]:
+                pixel_values_list.append(mm_content["image"]["pixel_values"])
+                image_grid_thw_list.append(mm_content["image"]["image_grid_thw"])
+            
+            # Process videos if present
+            if "video" in mm_content and mm_content["video"]:
+                pixel_values_videos_list.append(mm_content["video"]["pixel_values_videos"])
+                video_grid_thw_list.append(mm_content["video"]["video_grid_thw"])
+        
+        # Concatenate tensors
+        mm_content_dict = {}
+        if pixel_values_list:
+            mm_content_dict["pixel_values"] = torch.cat(pixel_values_list, dim=0) if len(pixel_values_list) > 1 else pixel_values_list[0]
+        if pixel_values_videos_list:
+            mm_content_dict["pixel_values_videos"] = torch.cat(pixel_values_videos_list, dim=0) if len(pixel_values_videos_list) > 1 else pixel_values_videos_list[0]
+        
+        # Prepare extra data
+        mm_extra_data = {}
+        if image_grid_thw_list:
+            mm_extra_data["image_grid_thw"] = torch.cat(image_grid_thw_list, dim=0) if len(image_grid_thw_list) > 1 else image_grid_thw_list[0]
+        if video_grid_thw_list:
+            mm_extra_data["video_grid_thw"] = torch.cat(video_grid_thw_list, dim=0) if len(video_grid_thw_list) > 1 else video_grid_thw_list[0]
+        
+        return mm_content_dict, mm_extra_data
 
+    @torch.inference_mode()
+    def forward(self, mm_data: List[Dict[str, Any]]):
+
+        mm_content_data, mm_extra_data = self._parse_and_batch_mm_data(mm_data)
+        pixel_values = mm_content_data.get("pixel_values", None)
+        pixel_values_videos = mm_content_data.get("pixel_values_videos", None)
+        
+        image_grid_thw = mm_extra_data.get("image_grid_thw", None)
+        video_grid_thw = mm_extra_data.get("video_grid_thw", None)
+        
+        embeds = []
+        if pixel_values is not None:
+            pixel_values = pixel_values.to(self.visual.dtype)
+            embeds.append(self.visual(pixel_values, grid_thw=image_grid_thw))
 
-class Qwen2_5_VLInputProcessor(Qwen2VLInputProcessorBase):
+        if pixel_values_videos is not None:
+            pixel_values_videos = pixel_values_videos.to(self.visual.dtype)
+            embeds.append(
+                self.visual(pixel_values_videos, grid_thw=video_grid_thw))
 
-    @classmethod
-    def get_model_class(cls):
-        return Qwen2_5_VLForConditionalGeneration
+        return embeds
 
 
 class Qwen2VLModelBase(PreTrainedModel):
@@ -413,7 +448,7 @@ def __init__(
         model_config.pretrained_config.rope_scaling['type'] = 'mrope'
         config = model_config.pretrained_config
 
-        assert model_config.attn_backend == 'TRTLLM', "Qwen2VL only supports TRTLLM backend now"
+        assert model_config.attn_backend == 'TRTLLM', "Qwen2/2.5-VL only supports TRTLLM backend now"
         super().__init__(config)
 
         self.model_config = model_config
@@ -458,26 +493,29 @@ def forward(
             f"num_context_requests: {num_context_requests}, num_generation_requests: {num_generation_requests}"
         )
 
-        mm_embed = kwargs.get("multi_modal_data", [])
-
-        error_msg = "Number of multimodal features (if provided) should be equal to number of context requests"
-        assert mm_embed == [] or len(
-            mm_embed) == num_context_requests, error_msg
-
-        input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
-                                                    input_ids, mm_embed)
+        mm_data = kwargs.get("mm_data", [])
+        mm_embeds = []
+        if len(mm_data) > 0:
+            assert len(
+                mm_data
+            ) == num_context_requests, f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch."
+            mm_embeds = self.mm_encoder.forward(mm_data)
 
         mrope_config = kwargs.get("mrope_config", {})
         if mrope_config:
             if mrope_rotary_cos_sin := mrope_config.get('mrope_rotary_cos_sin'):
+                assert len(mrope_rotary_cos_sin) == num_context_requests, f"Number of mrope_rotary_cos_sin ({len(mrope_rotary_cos_sin)}) should be equal to number of context requests ({num_context_requests}) in the batch."
                 mrope_config['mrope_rotary_cos_sin'] = torch.cat(
                     mrope_rotary_cos_sin, dim=0)
 
             if mrope_position_deltas := mrope_config.get(
                     'mrope_position_deltas'):
+                assert len(mrope_position_deltas) == num_generation_requests, f"Number of mrope_position_deltas ({len(mrope_position_deltas)}) should be equal to number of generation requests ({num_generation_requests}) in the batch."
                 mrope_config['mrope_position_deltas'] = torch.cat(
                     mrope_position_deltas, dim=0)
 
+        input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
+                                            input_ids, mm_embeds)
         output_prob = self.llm.forward(
             attn_metadata=attn_metadata,
             input_ids=input_ids,
@@ -485,17 +523,21 @@ def forward(
             inputs_embeds=input_embeds,
             return_context_logits=return_context_logits,
             mrope_config=mrope_config)
+
         logger.debug(f'output shape: {output_prob.shape}')
         return output_prob
 
 
 @register_auto_model("Qwen2VLForConditionalGeneration")
-@register_input_processor(Qwen2VLInputProcessor, model_type="qwen2_vl")
+@register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_vl")
 class Qwen2VLModel(Qwen2VLModelBase):
-    pass
-
+    def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs):
+        self.mm_encoder = Qwen2VisionModelBase(model_config.pretrained_config, Qwen2VLForConditionalGeneration)
+        super().__init__(model_config, *args, **kwargs)
 
 @register_auto_model("Qwen2_5_VLForConditionalGeneration")
-@register_input_processor(Qwen2_5_VLInputProcessor, model_type="qwen2_5_vl")
+@register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_5_vl")
 class Qwen2_5_VLModel(Qwen2VLModelBase):
-    pass
+    def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs):
+        super().__init__(model_config, *args, **kwargs)
+        self.mm_encoder = Qwen2VisionModelBase(model_config.pretrained_config, Qwen2_5_VLForConditionalGeneration)
\ No newline at end of file
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index 8e33920f6ef5..1c7c30873e60 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -281,6 +281,8 @@ def __init__(
             **kwargs):
         self.py_logits_post_processors = kwargs.pop("py_logits_post_processors",
                                                     None)
+        # Multimodal data
+        self.py_mm_data = kwargs.pop("py_mm_data", None)
         super().__init__(
             *args,
             client_id=client_id,
@@ -460,6 +462,7 @@ def executor_request_to_llm_request(
         if executor_request.client_id is not None else req_id,
         priority=0.5,
         llm_request_type=llm_request_type,
-        context_phase_params=executor_request.context_phase_params)
-
+        context_phase_params=executor_request.context_phase_params,
+        py_mm_data=getattr(executor_request, "py_mm_data",
+                                     None))
     return llm_request
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 7603e724c70a..fac1f96ceb4b 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1495,7 +1495,8 @@ def previous_seq_slots_device():
             self.position_ids_cuda[:total_num_tokens].unsqueeze(0),
             'inputs_embeds': None,
             'multi_modal_data': multi_modal_data,
-            'mrope_config': mrope_config
+            'mrope_config': mrope_config,
+            'mm_data': py_mm_data
         }
 
         if bool(lora_params):
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index ab501c85be53..c2251e6d50c3 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -1158,8 +1158,8 @@ def _update_new_active_requests_queue_latency(
     def _broadcast_new_requests(
         self,
         new_requests: List[RequestQueueItem],
-        py_request_objects: Optional[tuple[str, dict]] = None,
-    ) -> tuple[List[RequestQueueItem], Optional[tuple[str, dict]]]:
+        py_request_objects: Optional[dict[str, tuple[str, dict]]] = None,
+    ) -> tuple[List[RequestQueueItem], Optional[dict[str, tuple[str, dict]]]]:
         """Broadcasts new_requests and optional Python-only metadata (`py_request_objects`) across pipeline stages.
            `py_request_objects` is a tuple of (attribute_name, {request_id: object}).
         """
@@ -1206,8 +1206,11 @@ def _fetch_new_requests(self) -> List[RequestQueueItem]:
                 total_max_num_active_requests - total_num_active_requests)
 
         if self.dist.rank == 0:
-            py_request_objects = self._collect_py_objects_from_requests(
+            py_logits_post_processors = self._collect_py_objects_from_requests(
                 new_requests, "py_logits_post_processors")
+            py_mm_data = self._collect_py_objects_from_requests(
+                new_requests, "py_mm_data")
+            py_request_objects = tuple(filter(None, [py_logits_post_processors, py_mm_data]))
         else:
             py_request_objects = None
 
@@ -1234,9 +1237,9 @@ def _fetch_new_requests(self) -> List[RequestQueueItem]:
 
         if py_request_objects and (self.dist.tp_size > 1
                                    or self.dist.has_pp) and self.dist.rank > 0:
-            attr_name, req_obj_dict = py_request_objects
-            self._attach_py_objects_to_requests(new_requests, attr_name,
-                                                req_obj_dict)
+            for attr_name, req_obj_dict in py_request_objects:
+                self._attach_py_objects_to_requests(new_requests, attr_name,
+                                                    req_obj_dict)
 
         if not self.enable_attention_dp:
             self._update_new_active_requests_queue_latency(new_requests)
diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py
index fe218ecdcd55..dd4cbbab65e8 100644
--- a/tensorrt_llm/executor/executor.py
+++ b/tensorrt_llm/executor/executor.py
@@ -122,7 +122,8 @@ def generate_async(
             mrope_config: Optional[dict] = None,
             kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
             disaggregated_params: Optional[DisaggregatedParams] = None,
-            postproc_params: Optional[PostprocParams] = None
+            postproc_params: Optional[PostprocParams] = None,
+            mm_data: Optional[dict] = None,
     ) -> GenerationResult:
         """Generate output for the given prompt token ids in the asynchronous mode.
         Asynchronous generation accepts single prompt only.
@@ -148,7 +149,8 @@ def generate_async(
                 multimodal_embedding=multimodal_embedding,
                 mrope_config=mrope_config,
                 kv_cache_retention_config=kv_cache_retention_config,
-                disaggregated_params=disaggregated_params))
+                disaggregated_params=disaggregated_params,
+                mm_data=mm_data))
         return result
 
     def generate(
diff --git a/tensorrt_llm/executor/request.py b/tensorrt_llm/executor/request.py
index 933c9f435109..3c88cea63762 100644
--- a/tensorrt_llm/executor/request.py
+++ b/tensorrt_llm/executor/request.py
@@ -88,6 +88,7 @@ def __init__(
         kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
         disaggregated_params: Optional[DisaggregatedParams] = None,
         postproc_params: Optional[PostprocParams] = None,
+        mm_data: Optional[dict] = None,
     ):
         if isinstance(prompt_token_ids, list):
             self.prompt_token_ids = prompt_token_ids
@@ -112,6 +113,7 @@ def __init__(
         self.kv_cache_retention_config = kv_cache_retention_config
         self.id: Optional[int] = None
         self.disaggregated_params = disaggregated_params
+        self.mm_data = mm_data
 
     def set_id(self, id):
         assert self.id is None, f"Request ID is already set: {self.id}"
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index eeed86283d7c..3d872006af6e 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -479,6 +479,10 @@ def _deduce_max_tokens(request: GenerationRequest,
                 context_phase_params=context_phase_params,
                 type=request_type)
 
+            if self._is_pytorch_backend:
+                if request.mm_data is not None:
+                    executor_request.py_mm_data = request.mm_data
+
             if self._is_pytorch_backend and request.sampling_params.logits_processor:
                 # For PyTorch backend, we attach logits processors as a dynamic Python attribute
                 # instead of using the C++ binding, since the latter will cause PyCapsule pickling issues.
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index 8a32a580af57..a8cf5543b23a 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -357,9 +357,11 @@ def generate_async(
                 sampling_params.add_special_tokens = False
 
         query_token_ids = None
+        # NOTE: Multimodal related data
         multimodal_input = None
         multimodal_embedding = None
         mrope_config = None
+        mm_data = None
         if "prompt_token_ids" in inputs:
             # TODO: if specify prompt_token_ids, the mm hashing is not supported yet
             prompt_token_ids = inputs['prompt_token_ids']
@@ -384,11 +386,13 @@ def generate_async(
             prompt = inputs['prompt']
             if extra_processed_inputs is not None:
                 query_token_ids = extra_processed_inputs.get('query_token_ids')
+                # NOTE: Multimodal related data
                 multimodal_embedding = extra_processed_inputs.get(
                     'mm_embedding')
                 mrope_config = extra_processed_inputs.get('mrope_config')
                 multimodal_input = extra_processed_inputs.get(
                     'multimodal_input')
+                mm_data = extra_processed_inputs.get('mm_data')
         else:
             raise TypeError(
                 f"The inputs must be type str or list of int, but got {type(inputs)}"
@@ -414,6 +418,7 @@ def generate_async(
             kv_cache_retention_config=kv_cache_retention_config,
             disaggregated_params=disaggregated_params,
             postproc_params=_postproc_params,
+            mm_data=mm_data,
         )
 
         return RequestOutput._from_generation_result(result, prompt,

From ca5f3c066d343da099047b841b0b56e6e498accf Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 26 Jun 2025 23:58:32 +0900
Subject: [PATCH 02/14] address pre-commit

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_hyperclovax.py     | 52 +++++-----
 .../_torch/models/modeling_qwen2vl.py         | 99 ++++++++++++-------
 tensorrt_llm/_torch/pyexecutor/llm_request.py |  3 +-
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  3 +-
 tensorrt_llm/executor/executor.py             | 29 +++---
 5 files changed, 113 insertions(+), 73 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_hyperclovax.py b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
index e70c363a5b64..ff7c71b7a2e9 100644
--- a/tensorrt_llm/_torch/models/modeling_hyperclovax.py
+++ b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
@@ -1,14 +1,11 @@
 import copy
 import math
-import os
 from itertools import chain
 from typing import Any, Dict, List, Optional, Tuple, Union
 
-import numpy as np
 import torch
 import torch.nn as nn
 import transformers
-from PIL import Image
 from transformers import (AutoProcessor, AutoTokenizer, PretrainedConfig,
                           PreTrainedModel)
 
@@ -22,6 +19,7 @@
 from .modeling_multimodal_utils import fuse_input_embeds
 from .modeling_utils import register_auto_model
 
+
 # Copied from HyperCLOVAX-SEED-Vision-Instruct-3B/modeling_hyperclovax.py
 def select_best_resolution(original_size: tuple,
                            possible_resolutions: list) -> tuple:
@@ -527,23 +525,28 @@ def __call__(
 
         mm_data = {}
         mm_data["image"] = {
-            "pixel_values": torch.stack(preprocessed_image['pixel_values'][0], dim=0), #TODO change the pixel_values into the Shared Tensor
-            "image_sizes": preprocessed_image.get('image_sizes', None),
-            "is_videos": preprocessed_image.get('is_videos', None),
-            "num_queries_vis_abstractors": preprocessed_image.get('num_queries_vis_abstractors', None),
-            "num_queries_vis_abstractors_slow": preprocessed_image.get('num_queries_vis_abstractors_slow', None),
-            "first_last_frames_slows": preprocessed_image.get('first_last_frames_slows', None),
-        }
-        return fused_input_ids.to(torch.int32).tolist(), {
-            "mm_data": mm_data
+            "pixel_values":
+            torch.stack(
+                preprocessed_image['pixel_values'][0],
+                dim=0),  #TODO change the pixel_values into the Shared Tensor
+            "image_sizes":
+            preprocessed_image.get('image_sizes', None),
+            "is_videos":
+            preprocessed_image.get('is_videos', None),
+            "num_queries_vis_abstractors":
+            preprocessed_image.get('num_queries_vis_abstractors', None),
+            "num_queries_vis_abstractors_slow":
+            preprocessed_image.get('num_queries_vis_abstractors_slow', None),
+            "first_last_frames_slows":
+            preprocessed_image.get('first_last_frames_slows', None),
         }
+        return fused_input_ids.to(torch.int32).tolist(), {"mm_data": mm_data}
 
 
 class HCXVisionModel:
 
-    def __init__(self,
-                 pretrained_config: PretrainedConfig):
-    
+    def __init__(self, pretrained_config: PretrainedConfig):
+
         self.pretrained_config = pretrained_config
         self.vision_config = self.pretrained_config.vision_config
 
@@ -614,8 +617,13 @@ def _preprocess(self, mm_data: List[Any]) -> Dict[str, List[Any]]:
             for key in preprocessed_image_list[0].keys()
         }
 
-    def _parse_and_batch_mm_data(self, mm_data: List[Dict[str, Any]]) -> Tuple[List[torch.Tensor], Dict[str, List[Any]]]:
-        pixel_values = [list(torch.unbind(data["image"]["pixel_values"], dim=0)) for data in mm_data]
+    def _parse_and_batch_mm_data(
+        self, mm_data: List[Dict[str, Any]]
+    ) -> Tuple[List[torch.Tensor], Dict[str, List[Any]]]:
+        pixel_values = [
+            list(torch.unbind(data["image"]["pixel_values"], dim=0))
+            for data in mm_data
+        ]
         mm_extra_data = {
             key: [d["image"][key][0] for d in mm_data]
             for key in mm_data[0]["image"].keys()
@@ -624,7 +632,7 @@ def _parse_and_batch_mm_data(self, mm_data: List[Dict[str, Any]]) -> Tuple[List[
 
     @torch.inference_mode()
     def forward(self, mm_data: List[Dict[str, Any]]):
-        
+
         pixel_values, mm_extra_data = self._parse_and_batch_mm_data(mm_data)
         pixel_values = self._to_device(pixel_values)
         image_sizes = mm_extra_data.get("image_sizes", None)
@@ -633,8 +641,8 @@ def forward(self, mm_data: List[Dict[str, Any]]):
             "num_queries_vis_abstractors", None)
         num_queries_vis_abstractors_slow = mm_extra_data.get(
             "num_queries_vis_abstractors_slow", None)
-        first_last_frames_slows = mm_extra_data.get(
-            "first_last_frames_slows", None)
+        first_last_frames_slows = mm_extra_data.get("first_last_frames_slows",
+                                                    None)
 
         len_pixel_values = [len(pixel_value) for pixel_value in pixel_values]
         concat_pixel_values = torch.cat(list(chain(*pixel_values)),
@@ -825,9 +833,9 @@ def forward(
         mm_data = kwargs.get("mm_data", [])
         mm_embeds = []
         if len(mm_data) > 0:
-            assert len(
+            assert len(mm_data) == num_context_requests == len(
                 mm_data
-            ) == num_context_requests == len(mm_data), f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch."
+            ), f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch."
             mm_embeds = self.mm_encoder.forward(mm_data)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index a77c0697a02e..11a5b9aadb7d 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -28,7 +28,7 @@ def __init__(self,
                  trust_remote_code: bool = True):
         self.model_config = model_config
         self.tokenizer = tokenizer
-        # TODO: change to True and also change the acoording test result
+        # TODO: change to True and also change the according test result
         self.use_fast = False
         self.device = 'cuda'
         self.processor = AutoProcessor.from_pretrained(
@@ -277,7 +277,7 @@ def _preprocess(self, text: dict[str, any], mm_data: dict[str, any],
                               **mm_processor_kwargs)
 
     def _postprocess(self, input_ids: torch.IntTensor) -> torch.IntTensor:
-        # NOTE: Qwen2-VL's input processor is doing all the work for fusing input_ids with mm_tokens. 
+        # NOTE: Qwen2-VL's input processor is doing all the work for fusing input_ids with mm_tokens.
         # So, we just replace mm_tokens with expanded out-of-vocab ids
         masks = (input_ids == self.model_config.image_token_id) | (
             input_ids == self.model_config.vision_token_id) | (
@@ -334,7 +334,7 @@ def __call__(
 
         processed_inputs = self._preprocess(text_prompt, mm_data,
                                             mm_processor_kwargs).to(self.device)
-        
+
         if not mm_data:
             fused_input_ids = processed_inputs['input_ids']
             return fused_input_ids.to(torch.int32).tolist(), {}
@@ -342,14 +342,18 @@ def __call__(
         pixel_values = processed_inputs.get('pixel_values', None)
         pixel_values_videos = processed_inputs.get('pixel_values_videos', None)
         assert pixel_values is not None or pixel_values_videos is not None, "No multimodal data found"
-        
+
         mm_data = {}
         if pixel_values is not None:
-            mm_data["image"] = {"pixel_values": pixel_values, 
-                                "image_grid_thw": processed_inputs.get('image_grid_thw')}
+            mm_data["image"] = {
+                "pixel_values": pixel_values,
+                "image_grid_thw": processed_inputs.get('image_grid_thw')
+            }
         if pixel_values_videos is not None:
-            mm_data["video"] = {"pixel_values_videos": pixel_values_videos, 
-                                "video_grid_thw": processed_inputs.get('video_grid_thw')}
+            mm_data["video"] = {
+                "pixel_values_videos": pixel_values_videos,
+                "video_grid_thw": processed_inputs.get('video_grid_thw')
+            }
 
         input_ids = processed_inputs['input_ids']
         # TODO: We can move this to the LLM-side.
@@ -366,10 +370,14 @@ def __call__(
             "mm_data": mm_data,
         }
 
+
 class Qwen2VisionModelBase:
-    def __init__(self, pretrained_config: PretrainedConfig, model_class: type[PreTrainedModel]):
+
+    def __init__(self, pretrained_config: PretrainedConfig,
+                 model_class: type[PreTrainedModel]):
         self.pretrained_config = pretrained_config
-        self.device = torch.device('cuda' if torch.cuda.is_available() else 'cpu')
+        self.device = torch.device(
+            'cuda' if torch.cuda.is_available() else 'cpu')
 
         model_path = self.pretrained_config._name_or_path
         # TODO: Change the model class to TRT-LLM's Qwen2VisionModel
@@ -379,39 +387,53 @@ def __init__(self, pretrained_config: PretrainedConfig, model_class: type[PreTra
             torch_dtype=self.pretrained_config.torch_dtype,
             attn_implementation='flash_attention_2').eval()
         self.visual = model.visual.to(self.device)
-    
-    def _parse_and_batch_mm_data(self, mm_data: List[Dict[str, Any]]) -> Tuple[Dict[str, Any], Dict[str, List[Any]]]:
+
+    def _parse_and_batch_mm_data(
+        self, mm_data: List[Dict[str, Any]]
+    ) -> Tuple[Dict[str, Any], Dict[str, List[Any]]]:
 
         pixel_values_list = []
         pixel_values_videos_list = []
         image_grid_thw_list = []
         video_grid_thw_list = []
-        
+
         for mm_content in mm_data:
             # Process images if present
             if "image" in mm_content and mm_content["image"]:
                 pixel_values_list.append(mm_content["image"]["pixel_values"])
-                image_grid_thw_list.append(mm_content["image"]["image_grid_thw"])
-            
+                image_grid_thw_list.append(
+                    mm_content["image"]["image_grid_thw"])
+
             # Process videos if present
             if "video" in mm_content and mm_content["video"]:
-                pixel_values_videos_list.append(mm_content["video"]["pixel_values_videos"])
-                video_grid_thw_list.append(mm_content["video"]["video_grid_thw"])
-        
+                pixel_values_videos_list.append(
+                    mm_content["video"]["pixel_values_videos"])
+                video_grid_thw_list.append(
+                    mm_content["video"]["video_grid_thw"])
+
         # Concatenate tensors
         mm_content_dict = {}
         if pixel_values_list:
-            mm_content_dict["pixel_values"] = torch.cat(pixel_values_list, dim=0) if len(pixel_values_list) > 1 else pixel_values_list[0]
+            mm_content_dict["pixel_values"] = torch.cat(
+                pixel_values_list,
+                dim=0) if len(pixel_values_list) > 1 else pixel_values_list[0]
         if pixel_values_videos_list:
-            mm_content_dict["pixel_values_videos"] = torch.cat(pixel_values_videos_list, dim=0) if len(pixel_values_videos_list) > 1 else pixel_values_videos_list[0]
-        
+            mm_content_dict["pixel_values_videos"] = torch.cat(
+                pixel_values_videos_list,
+                dim=0) if len(pixel_values_videos_list
+                              ) > 1 else pixel_values_videos_list[0]
+
         # Prepare extra data
         mm_extra_data = {}
         if image_grid_thw_list:
-            mm_extra_data["image_grid_thw"] = torch.cat(image_grid_thw_list, dim=0) if len(image_grid_thw_list) > 1 else image_grid_thw_list[0]
+            mm_extra_data["image_grid_thw"] = torch.cat(
+                image_grid_thw_list, dim=0) if len(
+                    image_grid_thw_list) > 1 else image_grid_thw_list[0]
         if video_grid_thw_list:
-            mm_extra_data["video_grid_thw"] = torch.cat(video_grid_thw_list, dim=0) if len(video_grid_thw_list) > 1 else video_grid_thw_list[0]
-        
+            mm_extra_data["video_grid_thw"] = torch.cat(
+                video_grid_thw_list, dim=0) if len(
+                    video_grid_thw_list) > 1 else video_grid_thw_list[0]
+
         return mm_content_dict, mm_extra_data
 
     @torch.inference_mode()
@@ -420,10 +442,10 @@ def forward(self, mm_data: List[Dict[str, Any]]):
         mm_content_data, mm_extra_data = self._parse_and_batch_mm_data(mm_data)
         pixel_values = mm_content_data.get("pixel_values", None)
         pixel_values_videos = mm_content_data.get("pixel_values_videos", None)
-        
+
         image_grid_thw = mm_extra_data.get("image_grid_thw", None)
         video_grid_thw = mm_extra_data.get("video_grid_thw", None)
-        
+
         embeds = []
         if pixel_values is not None:
             pixel_values = pixel_values.to(self.visual.dtype)
@@ -504,18 +526,22 @@ def forward(
         mrope_config = kwargs.get("mrope_config", {})
         if mrope_config:
             if mrope_rotary_cos_sin := mrope_config.get('mrope_rotary_cos_sin'):
-                assert len(mrope_rotary_cos_sin) == num_context_requests, f"Number of mrope_rotary_cos_sin ({len(mrope_rotary_cos_sin)}) should be equal to number of context requests ({num_context_requests}) in the batch."
+                assert len(
+                    mrope_rotary_cos_sin
+                ) == num_context_requests, f"Number of mrope_rotary_cos_sin ({len(mrope_rotary_cos_sin)}) should be equal to number of context requests ({num_context_requests}) in the batch."
                 mrope_config['mrope_rotary_cos_sin'] = torch.cat(
                     mrope_rotary_cos_sin, dim=0)
 
             if mrope_position_deltas := mrope_config.get(
                     'mrope_position_deltas'):
-                assert len(mrope_position_deltas) == num_generation_requests, f"Number of mrope_position_deltas ({len(mrope_position_deltas)}) should be equal to number of generation requests ({num_generation_requests}) in the batch."
+                assert len(
+                    mrope_position_deltas
+                ) == num_generation_requests, f"Number of mrope_position_deltas ({len(mrope_position_deltas)}) should be equal to number of generation requests ({num_generation_requests}) in the batch."
                 mrope_config['mrope_position_deltas'] = torch.cat(
                     mrope_position_deltas, dim=0)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
-                                            input_ids, mm_embeds)
+                                                    input_ids, mm_embeds)
         output_prob = self.llm.forward(
             attn_metadata=attn_metadata,
             input_ids=input_ids,
@@ -531,13 +557,20 @@ def forward(
 @register_auto_model("Qwen2VLForConditionalGeneration")
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_vl")
 class Qwen2VLModel(Qwen2VLModelBase):
-    def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs):
-        self.mm_encoder = Qwen2VisionModelBase(model_config.pretrained_config, Qwen2VLForConditionalGeneration)
+
+    def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
+                 **kwargs):
+        self.mm_encoder = Qwen2VisionModelBase(model_config.pretrained_config,
+                                               Qwen2VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
 
+
 @register_auto_model("Qwen2_5_VLForConditionalGeneration")
 @register_input_processor(Qwen2VLInputProcessorBase, model_type="qwen2_5_vl")
 class Qwen2_5_VLModel(Qwen2VLModelBase):
-    def __init__(self, model_config: ModelConfig[PretrainedConfig], *args, **kwargs):
+
+    def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
+                 **kwargs):
         super().__init__(model_config, *args, **kwargs)
-        self.mm_encoder = Qwen2VisionModelBase(model_config.pretrained_config, Qwen2_5_VLForConditionalGeneration)
\ No newline at end of file
+        self.mm_encoder = Qwen2VisionModelBase(
+            model_config.pretrained_config, Qwen2_5_VLForConditionalGeneration)
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index 1c7c30873e60..e5c7ef250ee4 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -463,6 +463,5 @@ def executor_request_to_llm_request(
         priority=0.5,
         llm_request_type=llm_request_type,
         context_phase_params=executor_request.context_phase_params,
-        py_mm_data=getattr(executor_request, "py_mm_data",
-                                     None))
+        py_mm_data=getattr(executor_request, "py_mm_data", None))
     return llm_request
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index c2251e6d50c3..263f44e1f58b 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -1210,7 +1210,8 @@ def _fetch_new_requests(self) -> List[RequestQueueItem]:
                 new_requests, "py_logits_post_processors")
             py_mm_data = self._collect_py_objects_from_requests(
                 new_requests, "py_mm_data")
-            py_request_objects = tuple(filter(None, [py_logits_post_processors, py_mm_data]))
+            py_request_objects = tuple(
+                filter(None, [py_logits_post_processors, py_mm_data]))
         else:
             py_request_objects = None
 
diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py
index dd4cbbab65e8..9d805b17d79a 100644
--- a/tensorrt_llm/executor/executor.py
+++ b/tensorrt_llm/executor/executor.py
@@ -109,21 +109,20 @@ def abort_request(self, request_id: int) -> None:
         pass
 
     def generate_async(
-            self,
-            prompt_token_ids: List[int],
-            sampling_params: SamplingParams,
-            query_token_ids: Optional[Union[torch.Tensor, np.ndarray,
-                                            list]] = None,
-            lora_request: Optional[LoRARequest] = None,
-            prompt_adapter_request: Optional[PromptAdapterRequest] = None,
-            streaming: bool = False,
-            multimodal_input: Optional[MultimodalInput] = None,
-            multimodal_embedding: Optional[list] = None,
-            mrope_config: Optional[dict] = None,
-            kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
-            disaggregated_params: Optional[DisaggregatedParams] = None,
-            postproc_params: Optional[PostprocParams] = None,
-            mm_data: Optional[dict] = None,
+        self,
+        prompt_token_ids: List[int],
+        sampling_params: SamplingParams,
+        query_token_ids: Optional[Union[torch.Tensor, np.ndarray, list]] = None,
+        lora_request: Optional[LoRARequest] = None,
+        prompt_adapter_request: Optional[PromptAdapterRequest] = None,
+        streaming: bool = False,
+        multimodal_input: Optional[MultimodalInput] = None,
+        multimodal_embedding: Optional[list] = None,
+        mrope_config: Optional[dict] = None,
+        kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
+        disaggregated_params: Optional[DisaggregatedParams] = None,
+        postproc_params: Optional[PostprocParams] = None,
+        mm_data: Optional[dict] = None,
     ) -> GenerationResult:
         """Generate output for the given prompt token ids in the asynchronous mode.
         Asynchronous generation accepts single prompt only.

From d882ef870fdf0365999c78ee1b6148e834c3ca86 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 27 Jun 2025 20:01:47 +0900
Subject: [PATCH 03/14] change HyperCLOVAX-Vision Encoder from HF to TRT-LLM

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_hyperclovax.py     | 277 +++++++++++++++---
 .../_torch/models/modeling_qwen2vl.py         |  32 +-
 2 files changed, 264 insertions(+), 45 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_hyperclovax.py b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
index ff7c71b7a2e9..907bb0226339 100644
--- a/tensorrt_llm/_torch/models/modeling_hyperclovax.py
+++ b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
@@ -1,13 +1,16 @@
 import copy
 import math
+from functools import partial
 from itertools import chain
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 import torch.nn as nn
-import transformers
-from transformers import (AutoProcessor, AutoTokenizer, PretrainedConfig,
-                          PreTrainedModel)
+from einops import rearrange
+from transformers import (AutoConfig, AutoModel, AutoProcessor, AutoTokenizer,
+                          PretrainedConfig, PreTrainedModel)
+from transformers.modeling_utils import load_sharded_checkpoint
+from transformers.models.auto import CONFIG_MAPPING
 
 from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt,
                        register_input_processor)
@@ -17,6 +20,7 @@
 from ..model_config import ModelConfig
 from .modeling_auto import AutoModelForCausalLM
 from .modeling_multimodal_utils import fuse_input_embeds
+from .modeling_siglip import SiglipVisionModel
 from .modeling_utils import register_auto_model
 
 
@@ -398,6 +402,169 @@ def determine_non_vision_query_lengths(input_ids: torch.LongTensor, pad_id: int,
     return non_vision_query_lengths
 
 
+# Copied from HyperCLOVAX-SEED-Vision-Instruct-3B/modeling_hyperclovax.py
+class HCXVisionCAbstractor(nn.Module):
+    """
+    This module is based on C-Abstractor, whose license is under apache-2.0.
+    You can check the original code at https://github.com/khanrc/honeybee/blob/main/honeybee/projectors/projectors.py
+    and we made necessary modifications.
+    """
+
+    def __init__(
+        self,
+        num_queries: int,
+        num_input_tokens: int,
+        encoder_hidden_size: int,
+        hidden_size: int,
+        output_hidden_size: int,
+        pos_emb: bool = True,
+        prenorm: bool = False,
+    ):
+        super().__init__()
+        self.num_input_tokens = num_input_tokens
+        self.output_hidden_size = output_hidden_size
+
+        # Positional embedding
+        if pos_emb:
+            self.pos_emb = torch.nn.Parameter(
+                torch.zeros(1, num_input_tokens, encoder_hidden_size))
+            self.pos_emb.data.normal_(mean=0.0, std=0.02)
+        else:
+            self.pos_emb = None
+
+        # (Optional) Pre-normalization layer
+        from timm.layers import LayerNorm
+        if prenorm:
+            self.prenorm = LayerNorm(encoder_hidden_size)
+        else:
+            self.prenorm = None
+
+        self.build_net(num_queries, encoder_hidden_size, hidden_size,
+                       output_hidden_size)
+        self.dtype = next(self.parameters()).dtype
+
+    def forward(
+        self,
+        x: torch.Tensor,
+        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
+        num_grids: Optional[List[int]] = None,
+    ) -> torch.Tensor:
+        """
+        Args:
+            x: (B, L, encoder_hidden_size) tensor from the visual backbone (e.g. CLIP visual encoder), including cls token.
+        """
+        if self.prenorm is not None:
+            x = self.prenorm(x)
+
+        if self.pos_emb is not None:
+            x = x + self.pos_emb
+
+        x = self._forward(
+            x,
+            num_queries_vis_abstractors=num_queries_vis_abstractors,
+            num_grids=num_grids,
+        )  # (B, L, output_hidden_size)
+
+        return x
+
+    def _forward(
+        self,
+        x: torch.Tensor,
+        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
+        num_grids: Optional[List[int]] = None,
+    ) -> torch.Tensor:
+
+        # x: [B, L, dim]
+        B, L, dim = x.shape
+        hw = int(L**0.5)
+        x = rearrange(x, "b (h w) d -> b d h w", h=hw, w=hw)
+
+        if num_queries_vis_abstractors is not None:
+            assert num_grids is not None
+            return self._forward_adaptive_num_query(
+                x, num_queries_vis_abstractors, num_grids)
+
+        x = self.net(x)
+        x = rearrange(x, "b d h w -> b (h w) d")
+        x = self.readout(x)
+        return x
+
+    def _forward_adaptive_num_query(
+        self,
+        x: torch.Tensor,
+        num_queries_vis_abstractors: Optional[List[List[int]]] = None,
+        num_grids: Optional[List[int]] = None,
+    ) -> List[torch.Tensor]:
+        # self.net is consisted by 3 layers (s1, sampler, s2)
+        assert len(self.net) == 3
+
+        x = self.net[0](x)  # s1
+        new_x = []
+        for i, num_queries in enumerate(num_queries_vis_abstractors):
+            hw = int(num_queries**0.5)
+            sampler = nn.AdaptiveAvgPool2d((hw, hw))
+            out = sampler(x[num_grids[i]:num_grids[i + 1], :])
+            out = self.net[2](out)  # s2
+
+            out = rearrange(out, "b d h w -> b (h w) d")
+            out = self.readout(out)
+
+            new_x.append(out)
+        return new_x
+
+    def build_net(
+        self,
+        n_queries: int,
+        encoder_hidden_size: int,
+        hidden_size: int,
+        output_hidden_size: int,
+        depth: int = 3,
+        mlp_depth: int = 2,
+    ):
+        assert (n_queries**0.5).is_integer(
+        ), f"n_queries must be square number. n_queries: {n_queries}"
+        hw = int(n_queries**0.5)
+        from timm.layers import LayerNorm2d
+        from timm.models.regnet import RegStage
+
+        # RegBlock = ResBlock + SE
+        RegBlock = partial(
+            RegStage,
+            stride=1,
+            dilation=1,
+            act_layer=nn.SiLU,
+            norm_layer=LayerNorm2d,
+        )
+
+        s1 = RegBlock(
+            depth,
+            encoder_hidden_size,
+            hidden_size,
+        )
+        sampler = nn.AdaptiveAvgPool2d((hw, hw))
+        s2 = RegBlock(
+            depth,
+            hidden_size,
+            hidden_size,
+        )
+
+        self.net = nn.Sequential(s1, sampler, s2)
+        self.readout = self.build_mlp(mlp_depth, hidden_size,
+                                      output_hidden_size)
+
+    def build_mlp(
+        self,
+        depth: int,
+        hidden_size: int,
+        output_hidden_size: int,
+    ):
+        layers = [nn.Linear(hidden_size, output_hidden_size)]
+        for _ in range(1, depth):
+            layers.append(nn.SiLU())
+            layers.append(nn.Linear(output_hidden_size, output_hidden_size))
+        return nn.Sequential(*layers)
+
+
 class HCXVisionInputProcessor(InputProcessor):
 
     def __init__(self,
@@ -526,9 +693,9 @@ def __call__(
         mm_data = {}
         mm_data["image"] = {
             "pixel_values":
-            torch.stack(
-                preprocessed_image['pixel_values'][0],
-                dim=0),  #TODO change the pixel_values into the Shared Tensor
+            torch.stack(preprocessed_image['pixel_values'][0], dim=0).to(
+                torch.bfloat16
+            ),  #TODO change the pixel_values into the Shared Tensor
             "image_sizes":
             preprocessed_image.get('image_sizes', None),
             "is_videos":
@@ -545,23 +712,62 @@ def __call__(
 
 class HCXVisionModel:
 
-    def __init__(self, pretrained_config: PretrainedConfig):
+    def __init__(self, model_config: ModelConfig[PretrainedConfig]):
 
-        self.pretrained_config = pretrained_config
+        self.pretrained_config = model_config.pretrained_config
         self.vision_config = self.pretrained_config.vision_config
 
         model_path = self.pretrained_config._name_or_path
-
-        # NOTE: There is no way of importing mm_projector, HCXVisionCAbstractor from HF. So, can not do the sharded_loading.
-        # NOTE: trust_rmemote_code can be removed once we change the model into TRT-LLM's format
-        model = transformers.AutoModelForCausalLM.from_pretrained(
-            model_path, trust_remote_code=True).eval()
-        self.device = 'cuda'
-
-        # TODO: Convert to TRT-LLM's SIGLIP
-        self.vision_model = model.vision_model.to(self.device)
-        self.mm_projector = model.mm_projector.to(self.device)
-        self.image_newline = model.image_newline.to(self.device)
+        self.device = f"cuda:{model_config.mapping.rank}"
+
+        hf_model_config = AutoConfig.from_pretrained(model_path,
+                                                     trust_remote_code=True)
+        vision_model_type = hf_model_config.vision_config["model_type"]
+        vision_config = CONFIG_MAPPING[vision_model_type](
+            **hf_model_config.vision_config)
+        self.dtype = vision_config.torch_dtype
+        module_dict = nn.ModuleDict({
+            "vision_model":
+            AutoModel.from_config(vision_config, trust_remote_code=True),
+            "mm_projector":
+            HCXVisionCAbstractor(
+                num_queries=hf_model_config.num_queries_vis_abstractor,
+                num_input_tokens=(vision_config.image_size //
+                                  vision_config.patch_size)**2,
+                encoder_hidden_size=vision_config.hidden_size,
+                hidden_size=vision_config.hidden_size,
+                output_hidden_size=hf_model_config.
+                language_config["hidden_size"],
+                pos_emb=hf_model_config.proj_pos_emb,
+                prenorm=hf_model_config.proj_prenorm,
+            ),
+        })
+
+        module_dict.register_parameter(
+            "image_newline",
+            nn.Parameter(
+                torch.empty(hf_model_config.language_config["hidden_size"])))
+
+        missing_keys, _ = load_sharded_checkpoint(module_dict,
+                                                  model_path,
+                                                  strict=False)
+        assert len(missing_keys) == 0, f"Missing keys: {missing_keys}"
+        hf_vision_model = module_dict["vision_model"].to(self.dtype)
+        hf_mm_projector = module_dict["mm_projector"].to(self.dtype).to(
+            self.device)
+        hf_image_newline = module_dict.image_newline.to(self.dtype).to(
+            self.device)
+
+        vision_model_config = ModelConfig(pretrained_config=vision_config,
+                                          attn_backend="TRTLLM")
+
+        # Model related lines
+        self.vision_model = SiglipVisionModel(vision_model_config).to(
+            self.device).to(self.dtype)
+        self.vision_model.load_weights(hf_vision_model.state_dict())
+        print(model_config.mapping.rank, self.vision_model)
+        self.mm_projector = hf_mm_projector.eval()
+        self.image_newline = hf_image_newline
 
         self.unpad = self.pretrained_config.unpad
         self.use_nth_layer = self.pretrained_config.use_nth_layer
@@ -587,9 +793,11 @@ def _init_possible_resolutions(self, config: PretrainedConfig):
         return possible_resolutions
 
     def _to_device(
-            self, input_tensor: Union[torch.Tensor,
-                                      List]) -> Union[torch.Tensor, List]:
-        if isinstance(input_tensor, list):
+        self, input_tensor: Union[torch.Tensor, List, None]
+    ) -> Union[torch.Tensor, List, None]:
+        if input_tensor is None:
+            return None
+        elif isinstance(input_tensor, list):
             return [self._to_device(item) for item in input_tensor]
         elif isinstance(input_tensor, torch.Tensor):
             return input_tensor.to(self.device)
@@ -634,7 +842,8 @@ def _parse_and_batch_mm_data(
     def forward(self, mm_data: List[Dict[str, Any]]):
 
         pixel_values, mm_extra_data = self._parse_and_batch_mm_data(mm_data)
-        pixel_values = self._to_device(pixel_values)
+        pixel_values = self._to_device(
+            pixel_values)  # TODO: remove this once we have the shared tensor
         image_sizes = mm_extra_data.get("image_sizes", None)
         is_videos = mm_extra_data.get("is_videos", None)
         num_queries_vis_abstractors = mm_extra_data.get(
@@ -669,15 +878,15 @@ def forward(self, mm_data: List[Dict[str, Any]]):
                     device=concat_pixel_values.device,
                 )
                 chunk = torch.cat([chunk, dummy], dim=0)
-
+            attn_metadata = self.vision_model.prepare_attn_metadata(
+                chunk.shape[0])
             if self.use_nth_layer == -1:
                 self.vision_model.vision_model.post_layernorm = nn.Identity()
-                outs = self.vision_model(chunk)
-                outs = outs.last_hidden_state[:, visual_token_idx:]
+                outs = self.vision_model(chunk, attn_metadata=attn_metadata)
+                outs = outs[:, visual_token_idx:]
             else:
-                outs = self.vision_model(chunk, output_hidden_states=True)
-                outs = outs.hidden_states[self.use_nth_layer][:,
-                                                              visual_token_idx:]
+                outs = self.vision_model(chunk, attn_metadata=attn_metadata)
+                outs = outs[self.use_nth_layer][:, visual_token_idx:]
             image_forward_outs_chunks.append(outs)
 
         image_forward_outs = torch.cat(image_forward_outs_chunks, dim=0).to(
@@ -689,8 +898,6 @@ def forward(self, mm_data: List[Dict[str, Any]]):
             if is_videos is not None:
                 is_videos = list(chain(*is_videos))
             group_ids = None
-            image_forward_outs = image_forward_outs.to(
-                dtype=self.mm_projector.dtype)
             image_forward_outs = self.mm_projector(image_forward_outs)
         else:
             (
@@ -707,9 +914,6 @@ def forward(self, mm_data: List[Dict[str, Any]]):
                 is_videos,
                 first_last_frames_slows,
             )
-
-            image_forward_outs = image_forward_outs.to(
-                dtype=self.mm_projector.dtype)
             image_forward_outs = self.mm_projector(
                 image_forward_outs,
                 num_queries_vis_abstractors=num_queries_vis_abstractors,
@@ -779,14 +983,13 @@ def __init__(self, model_config: ModelConfig):
         if hasattr(self, "llm"):
             return
 
-        self.mm_encoder = HCXVisionModel(model_config.pretrained_config)
-
+        self.mm_encoder = HCXVisionModel(model_config)
         llm_model_config = copy.deepcopy(model_config)
         llm_model_config.pretrained_config = PretrainedConfig.from_dict(
             llm_model_config.pretrained_config.language_config)
         self.llm = AutoModelForCausalLM.from_config(llm_model_config)
 
-        self.model_dtype = getattr(config, "torch_dtype", torch.float16)
+        self.model_dtype = getattr(config, "torch_dtype", torch.bfloat16)
         logger.info(f"{self.dtype=} {self.model_dtype=}")
         self.post_config()
         self.is_loaded = True
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 11a5b9aadb7d..30c656c5490a 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -1,5 +1,5 @@
 import copy
-from typing import Any, Dict, List, Optional, Tuple
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
 from transformers import (AutoProcessor, AutoTokenizer, PretrainedConfig,
@@ -373,14 +373,14 @@ def __call__(
 
 class Qwen2VisionModelBase:
 
-    def __init__(self, pretrained_config: PretrainedConfig,
+    def __init__(self, model_config: ModelConfig[PretrainedConfig],
                  model_class: type[PreTrainedModel]):
-        self.pretrained_config = pretrained_config
-        self.device = torch.device(
-            'cuda' if torch.cuda.is_available() else 'cpu')
+        self.pretrained_config = model_config.pretrained_config
+        self.device = f"cuda:{model_config.mapping.rank}"
 
         model_path = self.pretrained_config._name_or_path
         # TODO: Change the model class to TRT-LLM's Qwen2VisionModel
+        # Currently, copying vision encoder on all devices.
         # NOTE: Using attn_implementation='flash_attention_2' to avoid the issue of vision model's GPU OOM.
         model = model_class.from_pretrained(
             model_path,
@@ -388,6 +388,16 @@ def __init__(self, pretrained_config: PretrainedConfig,
             attn_implementation='flash_attention_2').eval()
         self.visual = model.visual.to(self.device)
 
+    def _to_device(
+        self, input_tensor: Union[torch.Tensor, List, None]
+    ) -> Union[torch.Tensor, List, None]:
+        if input_tensor is None:
+            return None
+        elif isinstance(input_tensor, list):
+            return [self._to_device(item) for item in input_tensor]
+        elif isinstance(input_tensor, torch.Tensor):
+            return input_tensor.to(self.device)
+
     def _parse_and_batch_mm_data(
         self, mm_data: List[Dict[str, Any]]
     ) -> Tuple[Dict[str, Any], Dict[str, List[Any]]]:
@@ -448,10 +458,16 @@ def forward(self, mm_data: List[Dict[str, Any]]):
 
         embeds = []
         if pixel_values is not None:
+            pixel_values = self._to_device(
+                pixel_values
+            )  # TODO: remove this once we have the shared tensor
+            image_grid_thw = self._to_device(image_grid_thw)
             pixel_values = pixel_values.to(self.visual.dtype)
             embeds.append(self.visual(pixel_values, grid_thw=image_grid_thw))
 
         if pixel_values_videos is not None:
+            pixel_values_videos = self._to_device(pixel_values_videos)
+            video_grid_thw = self._to_device(video_grid_thw)
             pixel_values_videos = pixel_values_videos.to(self.visual.dtype)
             embeds.append(
                 self.visual(pixel_values_videos, grid_thw=video_grid_thw))
@@ -560,7 +576,7 @@ class Qwen2VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
-        self.mm_encoder = Qwen2VisionModelBase(model_config.pretrained_config,
+        self.mm_encoder = Qwen2VisionModelBase(model_config,
                                                Qwen2VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
 
@@ -571,6 +587,6 @@ class Qwen2_5_VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
-        super().__init__(model_config, *args, **kwargs)
         self.mm_encoder = Qwen2VisionModelBase(
-            model_config.pretrained_config, Qwen2_5_VLForConditionalGeneration)
+            model_config, Qwen2_5_VLForConditionalGeneration)
+        super().__init__(model_config, *args, **kwargs)

From 24499b44f2c8d2f1cfa1910093f5ffe84169375a Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Mon, 30 Jun 2025 20:11:26 +0900
Subject: [PATCH 04/14] add MultimodalParams

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_hyperclovax.py     | 69 ++++++--------
 .../_torch/models/modeling_qwen2vl.py         | 93 +++++++++++--------
 tensorrt_llm/_torch/pyexecutor/llm_request.py | 39 +++++---
 .../_torch/pyexecutor/model_engine.py         | 67 ++++++++-----
 tensorrt_llm/_torch/pyexecutor/py_executor.py |  6 +-
 tensorrt_llm/executor/executor.py             | 12 +--
 tensorrt_llm/executor/request.py              | 12 +--
 tensorrt_llm/executor/worker.py               | 48 ++++++----
 tensorrt_llm/inputs/multimodal.py             | 43 ++++++++-
 tensorrt_llm/llmapi/llm.py                    | 32 +++----
 10 files changed, 252 insertions(+), 169 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_hyperclovax.py b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
index 907bb0226339..dba63a91ead5 100644
--- a/tensorrt_llm/_torch/models/modeling_hyperclovax.py
+++ b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
@@ -12,6 +12,8 @@
 from transformers.modeling_utils import load_sharded_checkpoint
 from transformers.models.auto import CONFIG_MAPPING
 
+from tensorrt_llm.inputs.multimodal import MultimodalParams
+
 from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt,
                        register_input_processor)
 from ...logger import logger
@@ -690,8 +692,8 @@ def __call__(
         if not preprocessed_image:
             return fused_input_ids.to(torch.int32).tolist(), {}
 
-        mm_data = {}
-        mm_data["image"] = {
+        multimodal_data = {}
+        multimodal_data["image"] = {
             "pixel_values":
             torch.stack(preprocessed_image['pixel_values'][0], dim=0).to(
                 torch.bfloat16
@@ -707,7 +709,9 @@ def __call__(
             "first_last_frames_slows":
             preprocessed_image.get('first_last_frames_slows', None),
         }
-        return fused_input_ids.to(torch.int32).tolist(), {"mm_data": mm_data}
+        return fused_input_ids.to(torch.int32).tolist(), {
+            "multimodal_data": multimodal_data
+        }
 
 
 class HCXVisionModel:
@@ -765,7 +769,6 @@ def __init__(self, model_config: ModelConfig[PretrainedConfig]):
         self.vision_model = SiglipVisionModel(vision_model_config).to(
             self.device).to(self.dtype)
         self.vision_model.load_weights(hf_vision_model.state_dict())
-        print(model_config.mapping.rank, self.vision_model)
         self.mm_projector = hf_mm_projector.eval()
         self.image_newline = hf_image_newline
 
@@ -802,46 +805,30 @@ def _to_device(
         elif isinstance(input_tensor, torch.Tensor):
             return input_tensor.to(self.device)
 
-    # TODO: Remove this when we refactor LlmRequuest
-    def _preprocess(self, mm_data: List[Any]) -> Dict[str, List[Any]]:
-        preprocessed_image_list = []
-
-        for images in mm_data:
-            images = torch.unbind(images, dim=0)
-            preprocessed_image = self.processor(
-                images=images,
-                is_video_list=[False] * len(images),
-            )
-
-            # NOTE: The HCXVisionInputProcessor makes pixel_vlues to CPU values even though use_fast = True.
-            # So, we need to transfer them to GPU.
-            preprocessed_image["pixel_values"] = self._to_device(
-                preprocessed_image["pixel_values"])
-
-            preprocessed_image_list.append(preprocessed_image)
-
-        return {
-            key: [d[key][0] for d in preprocessed_image_list]
-            for key in preprocessed_image_list[0].keys()
-        }
-
-    def _parse_and_batch_mm_data(
-        self, mm_data: List[Dict[str, Any]]
+    def _parse_and_batch_multimodal_data(
+        self, multimodal_params: List[MultimodalParams]
     ) -> Tuple[List[torch.Tensor], Dict[str, List[Any]]]:
+        """Parse and batch multimodal data from MultimodalParams objects."""
         pixel_values = [
-            list(torch.unbind(data["image"]["pixel_values"], dim=0))
-            for data in mm_data
+            list(
+                torch.unbind(
+                    multimodal_param.multimodal_data["image"]["pixel_values"],
+                    dim=0)) for multimodal_param in multimodal_params
         ]
         mm_extra_data = {
-            key: [d["image"][key][0] for d in mm_data]
-            for key in mm_data[0]["image"].keys()
+            key: [
+                multimodal_param.multimodal_data["image"][key][0]
+                for multimodal_param in multimodal_params
+            ]
+            for key in multimodal_params[0].multimodal_data["image"].keys()
         }
         return pixel_values, mm_extra_data
 
     @torch.inference_mode()
-    def forward(self, mm_data: List[Dict[str, Any]]):
+    def forward(self, multimodal_params: List[MultimodalParams]):
 
-        pixel_values, mm_extra_data = self._parse_and_batch_mm_data(mm_data)
+        pixel_values, mm_extra_data = self._parse_and_batch_multimodal_data(
+            multimodal_params)
         pixel_values = self._to_device(
             pixel_values)  # TODO: remove this once we have the shared tensor
         image_sizes = mm_extra_data.get("image_sizes", None)
@@ -1033,13 +1020,13 @@ def forward(
             f"num_context_requests: {num_context_requests}, num_generation_requests: {num_generation_requests}"
         )
 
-        mm_data = kwargs.get("mm_data", [])
+        multimodal_params = kwargs.get("multimodal_params", [])
         mm_embeds = []
-        if len(mm_data) > 0:
-            assert len(mm_data) == num_context_requests == len(
-                mm_data
-            ), f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch."
-            mm_embeds = self.mm_encoder.forward(mm_data)
+        if len(multimodal_params) > 0:
+            assert len(multimodal_params) == num_context_requests == len(
+                multimodal_params
+            ), f"Number of multimodal tensors ({len(multimodal_params)}) should be equal to number of context requests ({num_context_requests}) in the batch."
+            mm_embeds = self.mm_encoder.forward(multimodal_params)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
                                                     input_ids, mm_embeds)
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 30c656c5490a..e2ce5635ad14 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -7,6 +7,8 @@
                           Qwen2VLForConditionalGeneration)
 from transformers.models.qwen2_vl.image_processing_qwen2_vl import smart_resize
 
+from tensorrt_llm.inputs.multimodal import MultimodalParams
+
 from ...functional import RopeEmbeddingUtils, RotaryScalingType
 from ...inputs import (ExtraProcessedInputs, InputProcessor, TextPrompt,
                        register_input_processor)
@@ -343,14 +345,14 @@ def __call__(
         pixel_values_videos = processed_inputs.get('pixel_values_videos', None)
         assert pixel_values is not None or pixel_values_videos is not None, "No multimodal data found"
 
-        mm_data = {}
+        multimodal_data = {}
         if pixel_values is not None:
-            mm_data["image"] = {
+            multimodal_data["image"] = {
                 "pixel_values": pixel_values,
                 "image_grid_thw": processed_inputs.get('image_grid_thw')
             }
         if pixel_values_videos is not None:
-            mm_data["video"] = {
+            multimodal_data["video"] = {
                 "pixel_values_videos": pixel_values_videos,
                 "video_grid_thw": processed_inputs.get('video_grid_thw')
             }
@@ -367,7 +369,7 @@ def __call__(
 
         return fused_input_ids.to(torch.int32).tolist(), {
             "mrope_config": mrope_config,
-            "mm_data": mm_data,
+            "multimodal_data": multimodal_data,
         }
 
 
@@ -398,8 +400,8 @@ def _to_device(
         elif isinstance(input_tensor, torch.Tensor):
             return input_tensor.to(self.device)
 
-    def _parse_and_batch_mm_data(
-        self, mm_data: List[Dict[str, Any]]
+    def _parse_and_batch_multimodal_data(
+        self, multimodal_params: List[MultimodalParams]
     ) -> Tuple[Dict[str, Any], Dict[str, List[Any]]]:
 
         pixel_values_list = []
@@ -407,19 +409,23 @@ def _parse_and_batch_mm_data(
         image_grid_thw_list = []
         video_grid_thw_list = []
 
-        for mm_content in mm_data:
+        for multimodal_param in multimodal_params:
             # Process images if present
-            if "image" in mm_content and mm_content["image"]:
-                pixel_values_list.append(mm_content["image"]["pixel_values"])
+            if "image" in multimodal_param.multimodal_data and multimodal_param.multimodal_data[
+                    "image"]:
+                pixel_values_list.append(
+                    multimodal_param.multimodal_data["image"]["pixel_values"])
                 image_grid_thw_list.append(
-                    mm_content["image"]["image_grid_thw"])
+                    multimodal_param.multimodal_data["image"]["image_grid_thw"])
 
             # Process videos if present
-            if "video" in mm_content and mm_content["video"]:
+            if "video" in multimodal_param.multimodal_data and multimodal_param.multimodal_data[
+                    "video"]:
                 pixel_values_videos_list.append(
-                    mm_content["video"]["pixel_values_videos"])
+                    multimodal_param.multimodal_data["video"]
+                    ["pixel_values_videos"])
                 video_grid_thw_list.append(
-                    mm_content["video"]["video_grid_thw"])
+                    multimodal_param.multimodal_data["video"]["video_grid_thw"])
 
         # Concatenate tensors
         mm_content_dict = {}
@@ -447,9 +453,10 @@ def _parse_and_batch_mm_data(
         return mm_content_dict, mm_extra_data
 
     @torch.inference_mode()
-    def forward(self, mm_data: List[Dict[str, Any]]):
+    def forward(self, multimodal_params: List[MultimodalParams]):
 
-        mm_content_data, mm_extra_data = self._parse_and_batch_mm_data(mm_data)
+        mm_content_data, mm_extra_data = self._parse_and_batch_multimodal_data(
+            multimodal_params)
         pixel_values = mm_content_data.get("pixel_values", None)
         pixel_values_videos = mm_content_data.get("pixel_values_videos", None)
 
@@ -513,6 +520,31 @@ def post_config(self):
         self.config = self.llm.config
         self.model_config.pretrained_config = self.llm.config
 
+    def _parse_mrope_config(
+            self, multimodal_params: List[MultimodalParams]
+    ) -> dict[str, torch.Tensor]:
+        mrope_config = {}
+        mrope_rotary_cos_sin_list = []
+        mrope_position_deltas_list = []
+
+        for multimodal_param in multimodal_params:
+            if hasattr(multimodal_param,
+                       'mrope_config') and multimodal_param.mrope_config:
+                if 'mrope_rotary_cos_sin' in multimodal_param.mrope_config:
+                    mrope_rotary_cos_sin_list.append(
+                        multimodal_param.mrope_config['mrope_rotary_cos_sin'])
+                if 'mrope_position_deltas' in multimodal_param.mrope_config:
+                    mrope_position_deltas_list.append(
+                        multimodal_param.mrope_config['mrope_position_deltas'])
+
+        if mrope_rotary_cos_sin_list:
+            mrope_config['mrope_rotary_cos_sin'] = torch.cat(
+                mrope_rotary_cos_sin_list, dim=0)
+
+        if mrope_position_deltas_list:
+            mrope_config['mrope_position_deltas'] = torch.cat(
+                mrope_position_deltas_list, dim=0)
+
     @torch.inference_mode()
     def forward(
         self,
@@ -531,30 +563,17 @@ def forward(
             f"num_context_requests: {num_context_requests}, num_generation_requests: {num_generation_requests}"
         )
 
-        mm_data = kwargs.get("mm_data", [])
+        multimodal_params = kwargs.get("multimodal_params", [])
         mm_embeds = []
-        if len(mm_data) > 0:
+        mrope_config = {}
+
+        if len(multimodal_params) > 0:
             assert len(
-                mm_data
-            ) == num_context_requests, f"Number of multimodal tensors ({len(mm_data)}) should be equal to number of context requests ({num_context_requests}) in the batch."
-            mm_embeds = self.mm_encoder.forward(mm_data)
-
-        mrope_config = kwargs.get("mrope_config", {})
-        if mrope_config:
-            if mrope_rotary_cos_sin := mrope_config.get('mrope_rotary_cos_sin'):
-                assert len(
-                    mrope_rotary_cos_sin
-                ) == num_context_requests, f"Number of mrope_rotary_cos_sin ({len(mrope_rotary_cos_sin)}) should be equal to number of context requests ({num_context_requests}) in the batch."
-                mrope_config['mrope_rotary_cos_sin'] = torch.cat(
-                    mrope_rotary_cos_sin, dim=0)
-
-            if mrope_position_deltas := mrope_config.get(
-                    'mrope_position_deltas'):
-                assert len(
-                    mrope_position_deltas
-                ) == num_generation_requests, f"Number of mrope_position_deltas ({len(mrope_position_deltas)}) should be equal to number of generation requests ({num_generation_requests}) in the batch."
-                mrope_config['mrope_position_deltas'] = torch.cat(
-                    mrope_position_deltas, dim=0)
+                multimodal_params
+            ) == num_context_requests, f"Number of multimodal tensors ({len(multimodal_params)}) should be equal to number of context requests ({num_context_requests}) in the batch."
+
+            mm_embeds = self.mm_encoder.forward(multimodal_params)
+            mrope_config = self._parse_mrope_config(multimodal_params)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
                                                     input_ids, mm_embeds)
diff --git a/tensorrt_llm/_torch/pyexecutor/llm_request.py b/tensorrt_llm/_torch/pyexecutor/llm_request.py
index e5c7ef250ee4..461c5de941e7 100644
--- a/tensorrt_llm/_torch/pyexecutor/llm_request.py
+++ b/tensorrt_llm/_torch/pyexecutor/llm_request.py
@@ -282,7 +282,7 @@ def __init__(
         self.py_logits_post_processors = kwargs.pop("py_logits_post_processors",
                                                     None)
         # Multimodal data
-        self.py_mm_data = kwargs.pop("py_mm_data", None)
+        self.py_multimodal_data = kwargs.pop("py_multimodal_data", None)
         super().__init__(
             *args,
             client_id=client_id,
@@ -402,6 +402,22 @@ def executor_request_to_llm_request(
     stop_words_list = convert_wordlist(
         executor_request.stop_words) if executor_request.stop_words else None
 
+    # Extract multimodal fields from executor request
+    multimodal_hashes = None
+    multimodal_positions = None
+    multimodal_lengths = None
+    if executor_request.multimodal_input is not None:
+        multimodal_hashes = executor_request.multimodal_input.multimodal_hashes
+        multimodal_positions = executor_request.multimodal_input.multimodal_positions
+        multimodal_lengths = executor_request.multimodal_input.multimodal_lengths
+
+    # Extract mrope fields
+    mrope_rotary_cos_sin = None
+    mrope_position_deltas = None
+    if executor_request.mrope_config is not None:
+        mrope_rotary_cos_sin = executor_request.mrope_config.mrope_rotary_cos_sin
+        mrope_position_deltas = executor_request.mrope_config.mrope_position_deltas
+
     llm_request = LlmRequest(
         request_id=req_id,
         max_new_tokens=executor_request.max_tokens,
@@ -421,24 +437,18 @@ def executor_request_to_llm_request(
         is None else executor_request.prompt_tuning_config.embedding_table,
         prompt_vocab_size=None if executor_request.prompt_tuning_config is None
         else executor_request.prompt_tuning_config.embedding_table.shape[0],
-        multimodal_hashes=None if executor_request.multimodal_input is None else
-        executor_request.multimodal_input.multimodal_hashes,
-        multimodal_positions=None if executor_request.multimodal_input is None
-        else executor_request.multimodal_input.multimodal_positions,
-        multimodal_lengths=None if executor_request.multimodal_input is None
-        else executor_request.multimodal_input.multimodal_lengths,
-        multimodal_embedding=None if executor_request.multimodal_embedding
-        is None else executor_request.multimodal_embedding,
+        multimodal_hashes=multimodal_hashes,
+        multimodal_positions=multimodal_positions,
+        multimodal_lengths=multimodal_lengths,
+        multimodal_embedding=executor_request.multimodal_embedding,
         lora_task_id=executor_request.lora_config.task_id
         if executor_request.lora_config is not None else None,
         lora_weights=executor_request.lora_config.weights
         if executor_request.lora_config is not None else None,
         lora_config=executor_request.lora_config.config
         if executor_request.lora_config is not None else None,
-        mrope_rotary_cos_sin=None if executor_request.mrope_config is None else
-        executor_request.mrope_config.mrope_rotary_cos_sin,
-        mrope_position_deltas=None if executor_request.mrope_config is None else
-        executor_request.mrope_config.mrope_position_deltas,
+        mrope_rotary_cos_sin=mrope_rotary_cos_sin,
+        mrope_position_deltas=mrope_position_deltas,
         lookahead_config=None,
         return_log_probs=executor_request.output_config.return_log_probs,
         return_context_logits=executor_request.output_config.
@@ -463,5 +473,6 @@ def executor_request_to_llm_request(
         priority=0.5,
         llm_request_type=llm_request_type,
         context_phase_params=executor_request.context_phase_params,
-        py_mm_data=getattr(executor_request, "py_mm_data", None))
+        py_multimodal_data=getattr(executor_request, "py_multimodal_data",
+                                   None))
     return llm_request
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index fac1f96ceb4b..07aa6ce716c4 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -10,7 +10,6 @@
 import traceback
 import weakref
 from abc import ABC, abstractmethod
-from collections import defaultdict
 from contextlib import contextmanager
 from typing import Any, Dict, List, Optional, Tuple
 
@@ -27,6 +26,7 @@
                                  local_mpi_size, nvtx_range, release_gc,
                                  torch_dtype_to_str, trace_func)
 from tensorrt_llm.bindings.executor import GuidedDecodingConfig
+from tensorrt_llm.inputs.multimodal import MultimodalParams
 from tensorrt_llm.logger import logger
 from tensorrt_llm.lora_manager import LoraConfig, LoraModelConfig
 from tensorrt_llm.mapping import Mapping
@@ -1197,19 +1197,38 @@ def _prepare_tp_inputs(
             prompt_lengths.append(len(prompt_tokens))
             past_seen_token_num = begin_compute
             num_cached_tokens_per_seq.append(past_seen_token_num)
-            multimodal_embedding = request.multimodal_embedding
-            if multimodal_embedding is not None:
-                multimodal_embedding = multimodal_embedding.pin_memory(
-                ) if multimodal_embedding.device == 'cpu' else multimodal_embedding
-                multi_modal_data.append(
-                    multimodal_embedding.to('cuda', non_blocking=True))
-
-            mrope_rotary_cos_sin = request.mrope_rotary_cos_sin
-            if mrope_rotary_cos_sin is not None:
-                mrope_rotary_cos_sin = mrope_rotary_cos_sin.pin_memory(
-                ) if mrope_rotary_cos_sin.device == 'cpu' else mrope_rotary_cos_sin
-                mrope_config['mrope_rotary_cos_sin'].append(
-                    mrope_rotary_cos_sin.to('cuda', non_blocking=True))
+
+            if request.multimodal_embedding is not None:
+                # TODO: Visit later once we have the SharedTensor.
+                request.multimodal_embedding = torch.tensor(
+                    request.multimodal_embedding,
+                    dtype=torch.float32,
+                    pin_memory=True)
+                request.multimodal_embedding = request.multimodal_embedding.to(
+                    'cuda', non_blocking=True)
+
+            if request.mrope_rotary_cos_sin is not None:
+                # TODO: Visit later once we have the SharedTensor.
+                mrope_rotary_cos_sin_tensor = torch.tensor(
+                    request.mrope_rotary_cos_sin,
+                    dtype=torch.float32,
+                    pin_memory=True)
+                mrope_rotary_cos_sin_tensor = mrope_rotary_cos_sin_tensor.to(
+                    'cuda', non_blocking=True)
+            else:
+                mrope_rotary_cos_sin_tensor = None
+            # Create MultimodalParams from request data
+            multimodal_params = MultimodalParams(
+                multimodal_embedding=request.multimodal_embedding,
+                mrope_config={
+                    'mrope_rotary_cos_sin': mrope_rotary_cos_sin_tensor
+                } if mrope_rotary_cos_sin_tensor is not None else None,
+                multimodal_data=request.py_multimodal_data,
+            )
+
+            if multimodal_params.has_content():
+                multimodal_params_list.append(multimodal_params)
+
             request.py_batch_idx = request.seq_slot
 
         num_ctx_requests = len(scheduled_requests.context_requests)
@@ -1238,13 +1257,19 @@ def _prepare_tp_inputs(
             else:
                 generation_requests.append(request)
 
+            # Handle generation request multimodal params
             mrope_position_deltas = request.mrope_position_deltas
             if mrope_position_deltas is not None:
-                mrope_position_deltas = torch.tensor([mrope_position_deltas],
-                                                     dtype=torch.int32,
-                                                     pin_memory=True)
-                mrope_config['mrope_position_deltas'].append(
-                    mrope_position_deltas.to('cuda', non_blocking=True))
+                mrope_position_deltas_tensor = torch.tensor(
+                    [mrope_position_deltas], dtype=torch.int32, pin_memory=True)
+                multimodal_params = MultimodalParams(
+                    mrope_config={
+                        'mrope_position_deltas':
+                        mrope_position_deltas_tensor.to('cuda',
+                                                        non_blocking=True)
+                    })
+                if multimodal_params.has_content():
+                    multimodal_params_list.append(multimodal_params)
         extend_requests += extend_dummy_requests
 
         if not self._disable_overlap_scheduler and self.is_spec_decode:
@@ -1494,9 +1519,7 @@ def previous_seq_slots_device():
             'position_ids':
             self.position_ids_cuda[:total_num_tokens].unsqueeze(0),
             'inputs_embeds': None,
-            'multi_modal_data': multi_modal_data,
-            'mrope_config': mrope_config,
-            'mm_data': py_mm_data
+            "multimodal_params": multimodal_params_list,
         }
 
         if bool(lora_params):
diff --git a/tensorrt_llm/_torch/pyexecutor/py_executor.py b/tensorrt_llm/_torch/pyexecutor/py_executor.py
index 263f44e1f58b..bceb97bb55e9 100644
--- a/tensorrt_llm/_torch/pyexecutor/py_executor.py
+++ b/tensorrt_llm/_torch/pyexecutor/py_executor.py
@@ -1208,10 +1208,10 @@ def _fetch_new_requests(self) -> List[RequestQueueItem]:
         if self.dist.rank == 0:
             py_logits_post_processors = self._collect_py_objects_from_requests(
                 new_requests, "py_logits_post_processors")
-            py_mm_data = self._collect_py_objects_from_requests(
-                new_requests, "py_mm_data")
+            py_multimodal_data = self._collect_py_objects_from_requests(
+                new_requests, "py_multimodal_data")
             py_request_objects = tuple(
-                filter(None, [py_logits_post_processors, py_mm_data]))
+                filter(None, [py_logits_post_processors, py_multimodal_data]))
         else:
             py_request_objects = None
 
diff --git a/tensorrt_llm/executor/executor.py b/tensorrt_llm/executor/executor.py
index 9d805b17d79a..ec4fc414cd2e 100644
--- a/tensorrt_llm/executor/executor.py
+++ b/tensorrt_llm/executor/executor.py
@@ -13,7 +13,7 @@
 import numpy as np
 import torch
 
-from tensorrt_llm.inputs.multimodal import MultimodalInput
+from tensorrt_llm.inputs.multimodal import MultimodalParams
 from tensorrt_llm.logger import logger, set_level
 from tensorrt_llm.lora_manager import LoraConfig
 
@@ -116,13 +116,10 @@ def generate_async(
         lora_request: Optional[LoRARequest] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
         streaming: bool = False,
-        multimodal_input: Optional[MultimodalInput] = None,
-        multimodal_embedding: Optional[list] = None,
-        mrope_config: Optional[dict] = None,
         kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
         disaggregated_params: Optional[DisaggregatedParams] = None,
         postproc_params: Optional[PostprocParams] = None,
-        mm_data: Optional[dict] = None,
+        multimodal_params: Optional[MultimodalParams] = None,
     ) -> GenerationResult:
         """Generate output for the given prompt token ids in the asynchronous mode.
         Asynchronous generation accepts single prompt only.
@@ -144,12 +141,9 @@ def generate_async(
                 lora_request=lora_request,
                 prompt_adapter_request=prompt_adapter_request,
                 streaming=streaming,
-                multimodal_input=multimodal_input,
-                multimodal_embedding=multimodal_embedding,
-                mrope_config=mrope_config,
                 kv_cache_retention_config=kv_cache_retention_config,
                 disaggregated_params=disaggregated_params,
-                mm_data=mm_data))
+                multimodal_params=multimodal_params))
         return result
 
     def generate(
diff --git a/tensorrt_llm/executor/request.py b/tensorrt_llm/executor/request.py
index 3c88cea63762..655af2610428 100644
--- a/tensorrt_llm/executor/request.py
+++ b/tensorrt_llm/executor/request.py
@@ -5,7 +5,7 @@
 import numpy as np
 import torch
 
-from tensorrt_llm.inputs.multimodal import MultimodalInput
+from tensorrt_llm.inputs.multimodal import MultimodalParams
 
 from ..disaggregated_params import DisaggregatedParams
 from ..llmapi.llm_utils import KvCacheRetentionConfig
@@ -82,13 +82,10 @@ def __init__(
         lora_request: Optional[LoRARequest] = None,
         prompt_adapter_request: Optional[PromptAdapterRequest] = None,
         streaming: bool = False,
-        multimodal_input: Optional[MultimodalInput] = None,
-        multimodal_embedding: Optional[list] = None,
-        mrope_config: Optional[dict] = None,
         kv_cache_retention_config: Optional[KvCacheRetentionConfig] = None,
         disaggregated_params: Optional[DisaggregatedParams] = None,
         postproc_params: Optional[PostprocParams] = None,
-        mm_data: Optional[dict] = None,
+        multimodal_params: Optional[MultimodalParams] = None,
     ):
         if isinstance(prompt_token_ids, list):
             self.prompt_token_ids = prompt_token_ids
@@ -107,13 +104,10 @@ def __init__(
         self.lora_request = lora_request
         self.prompt_adapter_request = prompt_adapter_request
         self.streaming = streaming
-        self.multimodal_input = multimodal_input
-        self.multimodal_embedding = multimodal_embedding
-        self.mrope_config = mrope_config
+        self.multimodal_params = multimodal_params
         self.kv_cache_retention_config = kv_cache_retention_config
         self.id: Optional[int] = None
         self.disaggregated_params = disaggregated_params
-        self.mm_data = mm_data
 
     def set_id(self, id):
         assert self.id is None, f"Request ID is already set: {self.id}"
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index 3d872006af6e..228410beae8a 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -376,11 +376,6 @@ def _enqueue_request(self, request: GenerationRequest) -> int:
 
         prompt_token_ids = copy.deepcopy(request.prompt_token_ids)
         prompt_tuning_config = None
-        multimodal_embedding = None
-        mrope_config = None
-        multimodal_input = None
-        if request.multimodal_embedding is not None:
-            multimodal_embedding = request.multimodal_embedding
         if request.prompt_adapter_request is not None:
             self._load_prompt_adapter(request.prompt_adapter_request)
             uid = str(request.prompt_adapter_request.adapter_id)
@@ -391,15 +386,35 @@ def _enqueue_request(self, request: GenerationRequest) -> int:
             prompt_token_ids = list(range(
                 vocab_size, vocab_size + pa_length)) + prompt_token_ids
 
-        if request.mrope_config is not None:
-            mrope_config = tllm.MropeConfig(**request.mrope_config)
-
-        if request.multimodal_input is not None:
-            multimodal_input = tllm.MultimodalInput(
-                multimodal_hashes=request.multimodal_input.multimodal_hashes,
-                multimodal_positions=request.multimodal_input.
-                multimodal_positions,
-                multimodal_lengths=request.multimodal_input.multimodal_lengths)
+        # Multimodal related fields - simplified handling
+        if request.multimodal_params is not None and request.multimodal_params.has_content(
+        ):
+            # Create mrope_config if needed
+            mrope_config = None
+            if request.multimodal_params.mrope_config:
+                mrope_config = tllm.MropeConfig(
+                    mrope_rotary_cos_sin=request.multimodal_params.mrope_config.
+                    get('mrope_rotary_cos_sin'),
+                    mrope_position_deltas=request.multimodal_params.
+                    mrope_config.get('mrope_position_deltas'))
+
+            # Create multimodal_input for C++ if needed
+            multimodal_input = None
+            if request.multimodal_params.multimodal_input is not None:
+                multimodal_input = tllm.MultimodalInput(
+                    multimodal_hashes=request.multimodal_params.
+                    multimodal_input.multimodal_hashes,
+                    multimodal_positions=request.multimodal_params.
+                    multimodal_input.multimodal_positions,
+                    multimodal_lengths=request.multimodal_params.
+                    multimodal_input.multimodal_lengths)
+            multimodal_embedding = None
+            if request.multimodal_params.multimodal_embedding is not None:
+                multimodal_embedding = request.multimodal_params.multimodal_embedding
+        else:
+            multimodal_embedding = None
+            mrope_config = None
+            multimodal_input = None
 
         context_phase_params = None
         request_type = tllm.RequestType.REQUEST_TYPE_CONTEXT_AND_GENERATION
@@ -480,8 +495,9 @@ def _deduce_max_tokens(request: GenerationRequest,
                 type=request_type)
 
             if self._is_pytorch_backend:
-                if request.mm_data is not None:
-                    executor_request.py_mm_data = request.mm_data
+                # For PyTorch backend, attach the raw multimodal data
+                if request.multimodal_params is not None and request.multimodal_params.multimodal_data:
+                    executor_request.py_multimodal_data = request.multimodal_params.multimodal_data
 
             if self._is_pytorch_backend and request.sampling_params.logits_processor:
                 # For PyTorch backend, we attach logits processors as a dynamic Python attribute
diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index 33d435ef1a7b..36f07d6486a8 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -1,7 +1,7 @@
 """Multimodal utilities for handling images and other media types in TensorRT-LLM."""
 
-from dataclasses import dataclass
-from typing import Any, Dict, List, Tuple, Union
+from dataclasses import dataclass, field
+from typing import Any, Dict, List, Optional, Tuple, Union
 
 import numpy as np
 import PIL
@@ -82,6 +82,45 @@ def to_tensor(self) -> Tuple[torch.Tensor, torch.Tensor, torch.Tensor]:
             torch.tensor(self.multimodal_lengths, dtype=torch.int32))
 
 
+@dataclass
+class MultimodalParams:
+    """Unified container for multimodal parameters.
+
+    This class encapsulates all multimodal-related data that flows through the system,
+    providing a clean interface for handling multimodal inputs across different models.
+    """
+
+    # Core multimodal data
+    multimodal_input: Optional[MultimodalInput] = None
+    """Multimodal input data with hashing information for caching and deduplication."""
+
+    multimodal_embedding: Optional[torch.Tensor] = None
+    """Pre-computed multimodal embeddings from vision encoder."""
+
+    multimodal_data: Optional[Dict[str, Dict[str,
+                                             Union[torch.Tensor,
+                                                   List[Any]]]]] = field(
+                                                       default_factory=dict)
+    """Raw multimodal data by modality (e.g., image pixels, video frames)."""
+
+    # Model-specific configurations
+    mrope_config: Optional[Dict[str, Any]] = None
+    """Multimodal rotary position embedding config (used by Qwen2-VL)."""
+
+    def __post_init__(self):
+        """Ensure default values are properly set."""
+        if self.multimodal_data is None:
+            self.multimodal_data = {}
+        if self.mrope_config is None:
+            self.mrope_config = {}
+
+    def has_content(self) -> bool:
+        """Check if this object contains any multimodal data."""
+        return bool(self.multimodal_input
+                    or self.multimodal_embedding is not None
+                    or self.multimodal_data or self.mrope_config)
+
+
 # adopt from vllm : https://github.com/vllm-project/vllm/blob/main/vllm/vllm/multimodal/hash.py
 def serialize_item(obj: object) -> bytes:
     # Simple cases
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index a8cf5543b23a..86d1be4b8882 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -13,6 +13,7 @@
 from transformers import PreTrainedTokenizerBase
 
 from tensorrt_llm.inputs.data import TextPrompt
+from tensorrt_llm.inputs.multimodal import MultimodalParams
 from tensorrt_llm.inputs.registry import DefaultInputProcessor
 
 from .._utils import nvtx_range_debug
@@ -357,11 +358,8 @@ def generate_async(
                 sampling_params.add_special_tokens = False
 
         query_token_ids = None
-        # NOTE: Multimodal related data
-        multimodal_input = None
-        multimodal_embedding = None
-        mrope_config = None
-        mm_data = None
+        multimodal_params = None
+
         if "prompt_token_ids" in inputs:
             # TODO: if specify prompt_token_ids, the mm hashing is not supported yet
             prompt_token_ids = inputs['prompt_token_ids']
@@ -386,13 +384,18 @@ def generate_async(
             prompt = inputs['prompt']
             if extra_processed_inputs is not None:
                 query_token_ids = extra_processed_inputs.get('query_token_ids')
-                # NOTE: Multimodal related data
-                multimodal_embedding = extra_processed_inputs.get(
-                    'mm_embedding')
-                mrope_config = extra_processed_inputs.get('mrope_config')
-                multimodal_input = extra_processed_inputs.get(
-                    'multimodal_input')
-                mm_data = extra_processed_inputs.get('mm_data')
+                # Create unified MultimodalParams
+                multimodal_params = MultimodalParams(
+                    multimodal_embedding=extra_processed_inputs.get(
+                        'mm_embedding'),
+                    mrope_config=extra_processed_inputs.get('mrope_config'),
+                    multimodal_input=extra_processed_inputs.get(
+                        'multimodal_input'),
+                    multimodal_data=extra_processed_inputs.get(
+                        'multimodal_data'))
+                # Only pass it if it has content
+                if not multimodal_params.has_content():
+                    multimodal_params = None
         else:
             raise TypeError(
                 f"The inputs must be type str or list of int, but got {type(inputs)}"
@@ -412,13 +415,10 @@ def generate_async(
             lora_request=lora_request,
             prompt_adapter_request=prompt_adapter_request,
             streaming=streaming,
-            multimodal_input=multimodal_input,
-            multimodal_embedding=multimodal_embedding,
-            mrope_config=mrope_config,
             kv_cache_retention_config=kv_cache_retention_config,
             disaggregated_params=disaggregated_params,
             postproc_params=_postproc_params,
-            mm_data=mm_data,
+            multimodal_params=multimodal_params,
         )
 
         return RequestOutput._from_generation_result(result, prompt,

From b8f76b7b2031c61081bfcf30853f7ffd7ff883a7 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 3 Jul 2025 19:08:36 +0900
Subject: [PATCH 05/14] multimodal_data into multimodal_params

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/models/modeling_gemma3vl.py   | 6 +++++-
 tensorrt_llm/_torch/models/modeling_llava_next.py | 6 +++++-
 tensorrt_llm/_torch/models/modeling_vila.py       | 6 +++++-
 3 files changed, 15 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py
index 229c30e40c99..4391189daa73 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py
@@ -161,7 +161,11 @@ def forward(
             f"[Gemma3Model::forward]{num_context_requests=}, {num_generation_requests=}"
         )
 
-        mm_embed = kwargs.get("multi_modal_data", [])
+        multimodal_params = kwargs.get("multimodal_params", [])
+        mm_embed = [
+            multimodal_param.mm_embedding
+            for multimodal_param in multimodal_params
+        ]
         assert mm_embed == [] or len(
             mm_embed
         ) == num_context_requests, "Number of multimodal features (if provided) should be equal to number of context requests"
diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index 967555dcb389..5fc2fbdfa27b 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -271,7 +271,11 @@ def forward(
         num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations
         logger.debug(f"{num_context_requests=}, {num_generation_requests=}")
 
-        mm_embed = kwargs.get("multi_modal_data", [])
+        multimodal_params = kwargs.get("multimodal_params", [])
+        mm_embed = [
+            multimodal_param.mm_embedding
+            for multimodal_param in multimodal_params
+        ]
         assert mm_embed == [] or len(
             mm_embed
         ) == num_context_requests, "Number of multimodal features (if provided) should be equal to number of context requests"
diff --git a/tensorrt_llm/_torch/models/modeling_vila.py b/tensorrt_llm/_torch/models/modeling_vila.py
index 217ec9c388ab..c2f2d7c79081 100644
--- a/tensorrt_llm/_torch/models/modeling_vila.py
+++ b/tensorrt_llm/_torch/models/modeling_vila.py
@@ -1161,7 +1161,11 @@ def forward(
         """
 
         num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations
-        mm_embed = kwargs.get("multi_modal_data", [])
+        multimodal_params = kwargs.get("multimodal_params", [])
+        mm_embed = [
+            multimodal_param.mm_embedding
+            for multimodal_param in multimodal_params
+        ]
 
         assert mm_embed == [] or len(
             mm_embed

From 105ad8fa39c7430931dbb88f9cefc6e4168a55dc Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 3 Jul 2025 20:32:43 +0900
Subject: [PATCH 06/14] address pin memory

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_gemma3vl.py        |  2 +-
 .../_torch/models/modeling_llava_next.py      |  2 +-
 .../_torch/models/modeling_qwen2vl.py         |  7 ++--
 tensorrt_llm/_torch/models/modeling_vila.py   |  2 +-
 .../_torch/pyexecutor/model_engine.py         | 32 ++++++++-----------
 5 files changed, 19 insertions(+), 26 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py
index 4391189daa73..85cb733c5dfd 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py
@@ -163,7 +163,7 @@ def forward(
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.mm_embedding
+            multimodal_param.multimodal_embedding
             for multimodal_param in multimodal_params
         ]
         assert mm_embed == [] or len(
diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index 5fc2fbdfa27b..15f7f5452d3e 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -273,7 +273,7 @@ def forward(
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.mm_embedding
+            multimodal_param.multimodal_embedding
             for multimodal_param in multimodal_params
         ]
         assert mm_embed == [] or len(
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index e2ce5635ad14..917752d12fad 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -568,11 +568,8 @@ def forward(
         mrope_config = {}
 
         if len(multimodal_params) > 0:
-            assert len(
-                multimodal_params
-            ) == num_context_requests, f"Number of multimodal tensors ({len(multimodal_params)}) should be equal to number of context requests ({num_context_requests}) in the batch."
-
-            mm_embeds = self.mm_encoder.forward(multimodal_params)
+            mm_embeds = self.mm_encoder.forward(
+                multimodal_params[:num_context_requests])
             mrope_config = self._parse_mrope_config(multimodal_params)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
diff --git a/tensorrt_llm/_torch/models/modeling_vila.py b/tensorrt_llm/_torch/models/modeling_vila.py
index c2f2d7c79081..3653f9c43fec 100644
--- a/tensorrt_llm/_torch/models/modeling_vila.py
+++ b/tensorrt_llm/_torch/models/modeling_vila.py
@@ -1163,7 +1163,7 @@ def forward(
         num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.mm_embedding
+            multimodal_param.multimodal_embedding
             for multimodal_param in multimodal_params
         ]
 
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 07aa6ce716c4..1eb49b64e1fc 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1198,31 +1198,27 @@ def _prepare_tp_inputs(
             past_seen_token_num = begin_compute
             num_cached_tokens_per_seq.append(past_seen_token_num)
 
-            if request.multimodal_embedding is not None:
+            multimodal_embedding = request.multimodal_embedding
+            if multimodal_embedding is not None:
                 # TODO: Visit later once we have the SharedTensor.
-                request.multimodal_embedding = torch.tensor(
-                    request.multimodal_embedding,
-                    dtype=torch.float32,
-                    pin_memory=True)
-                request.multimodal_embedding = request.multimodal_embedding.to(
+                multimodal_embedding = multimodal_embedding.pin_memory(
+                ) if multimodal_embedding.device == 'cpu' else multimodal_embedding
+                multimodal_embedding = multimodal_embedding.to(
                     'cuda', non_blocking=True)
 
-            if request.mrope_rotary_cos_sin is not None:
+            mrope_rotary_cos_sin = request.mrope_rotary_cos_sin
+            if mrope_rotary_cos_sin is not None:
                 # TODO: Visit later once we have the SharedTensor.
-                mrope_rotary_cos_sin_tensor = torch.tensor(
-                    request.mrope_rotary_cos_sin,
-                    dtype=torch.float32,
-                    pin_memory=True)
-                mrope_rotary_cos_sin_tensor = mrope_rotary_cos_sin_tensor.to(
+                mrope_rotary_cos_sin = mrope_rotary_cos_sin.pin_memory(
+                ) if mrope_rotary_cos_sin.device == 'cpu' else mrope_rotary_cos_sin
+                mrope_rotary_cos_sin = mrope_rotary_cos_sin.to(
                     'cuda', non_blocking=True)
-            else:
-                mrope_rotary_cos_sin_tensor = None
+
             # Create MultimodalParams from request data
             multimodal_params = MultimodalParams(
-                multimodal_embedding=request.multimodal_embedding,
-                mrope_config={
-                    'mrope_rotary_cos_sin': mrope_rotary_cos_sin_tensor
-                } if mrope_rotary_cos_sin_tensor is not None else None,
+                multimodal_embedding=multimodal_embedding,
+                mrope_config={'mrope_rotary_cos_sin': mrope_rotary_cos_sin}
+                if mrope_rotary_cos_sin is not None else {},
                 multimodal_data=request.py_multimodal_data,
             )
 

From 9602cdb4c7738fd1ec92e945bf98138aa0ae3ba1 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 3 Jul 2025 23:35:40 +0900
Subject: [PATCH 07/14] add example of multimodal_data and pre-commit

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 tensorrt_llm/inputs/multimodal.py | 17 ++++++++++++++++-
 1 file changed, 16 insertions(+), 1 deletion(-)

diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index 36f07d6486a8..8d8b7962597a 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -101,7 +101,22 @@ class MultimodalParams:
                                              Union[torch.Tensor,
                                                    List[Any]]]]] = field(
                                                        default_factory=dict)
-    """Raw multimodal data by modality (e.g., image pixels, video frames)."""
+    """Processed multimodal data after AutoProcessor's process() by modality (e.g., image pixels, video pixel values).
+    It should be in the form of {modality: {item_str: item_data}}
+    e.g.
+    {
+        "image": {
+            "pixel_values": torch.Tensor(),
+            "image_height": torch.Tensor() or List[int],
+            "image_width": torch.Tensor() or List[int]
+        },
+        "video": {
+            "pixel_values": torch.Tensor(),
+            "video_height": torch.Tensor() or List[int],
+            "video_width": torch.Tensor() or List[int]
+        },
+    }
+    """
 
     # Model-specific configurations
     mrope_config: Optional[Dict[str, Any]] = None

From e2b57d754c4c17a7c94a8300709c06b192af948f Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 4 Jul 2025 15:16:15 +0900
Subject: [PATCH 08/14] move mrope_config and mm_embedding under
 MultimodalParams

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_gemma3vl.py        |  6 +-
 .../_torch/models/modeling_hyperclovax.py     | 15 ++-
 tensorrt_llm/_torch/models/modeling_llama.py  | 13 ++-
 .../_torch/models/modeling_llava_next.py      |  6 +-
 .../_torch/models/modeling_qwen2vl.py         | 56 +++++++----
 tensorrt_llm/_torch/models/modeling_vila.py   |  6 +-
 .../_torch/pyexecutor/model_engine.py         | 52 ++++-------
 tensorrt_llm/executor/worker.py               | 34 ++-----
 tensorrt_llm/inputs/multimodal.py             | 92 +++++++++++++++----
 tensorrt_llm/llmapi/llm.py                    |  3 -
 10 files changed, 169 insertions(+), 114 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_gemma3vl.py b/tensorrt_llm/_torch/models/modeling_gemma3vl.py
index 85cb733c5dfd..ceb6b01a9cea 100644
--- a/tensorrt_llm/_torch/models/modeling_gemma3vl.py
+++ b/tensorrt_llm/_torch/models/modeling_gemma3vl.py
@@ -100,8 +100,10 @@ def __call__(
             "pixel_values"]
         input_ids = preprocess_outputs[0]["mm_processor_kwargs"]["input_ids"]
         mm_features = self._process(pixel_values)
+        multimodal_data = {}
+        multimodal_data["multimodal_embedding"] = mm_features
         return input_ids[0].to(torch.int32).tolist(), {
-            "mm_embedding": mm_features
+            "multimodal_data": multimodal_data
         }
 
 
@@ -163,7 +165,7 @@ def forward(
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.multimodal_embedding
+            multimodal_param.multimodal_data["multimodal_embedding"]
             for multimodal_param in multimodal_params
         ]
         assert mm_embed == [] or len(
diff --git a/tensorrt_llm/_torch/models/modeling_hyperclovax.py b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
index dba63a91ead5..9f37759ba03b 100644
--- a/tensorrt_llm/_torch/models/modeling_hyperclovax.py
+++ b/tensorrt_llm/_torch/models/modeling_hyperclovax.py
@@ -1,5 +1,6 @@
 import copy
 import math
+import os
 from functools import partial
 from itertools import chain
 from typing import Any, Dict, List, Optional, Tuple, Union
@@ -25,6 +26,8 @@
 from .modeling_siglip import SiglipVisionModel
 from .modeling_utils import register_auto_model
 
+DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
+
 
 # Copied from HyperCLOVAX-SEED-Vision-Instruct-3B/modeling_hyperclovax.py
 def select_best_resolution(original_size: tuple,
@@ -969,8 +972,8 @@ def __init__(self, model_config: ModelConfig):
         self.model_config = model_config
         if hasattr(self, "llm"):
             return
-
-        self.mm_encoder = HCXVisionModel(model_config)
+        if not DISAGG:
+            self.mm_encoder = HCXVisionModel(model_config)
         llm_model_config = copy.deepcopy(model_config)
         llm_model_config.pretrained_config = PretrainedConfig.from_dict(
             llm_model_config.pretrained_config.language_config)
@@ -1026,7 +1029,13 @@ def forward(
             assert len(multimodal_params) == num_context_requests == len(
                 multimodal_params
             ), f"Number of multimodal tensors ({len(multimodal_params)}) should be equal to number of context requests ({num_context_requests}) in the batch."
-            mm_embeds = self.mm_encoder.forward(multimodal_params)
+            if not DISAGG:
+                mm_embeds = self.mm_encoder.forward(multimodal_params)
+            else:
+                mm_embeds = [
+                    multimodal_param.multimodal_data["multimodal_embedding"]
+                    for multimodal_param in multimodal_params
+                ]
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
                                                     input_ids, mm_embeds)
diff --git a/tensorrt_llm/_torch/models/modeling_llama.py b/tensorrt_llm/_torch/models/modeling_llama.py
index 8bd0cd80e6c5..8f2e26966091 100644
--- a/tensorrt_llm/_torch/models/modeling_llama.py
+++ b/tensorrt_llm/_torch/models/modeling_llama.py
@@ -851,7 +851,10 @@ def __call__(
             mm_embeds = self.encoder.multi_modal_projector(mm_embeds)
             # for fuse_input_embeds
             token_ids[token_ids == self.image_token_index] = self.vocab_size + 1
-            return token_ids.tolist(), {"mm_embedding": mm_embeds}
+
+            multimodal_data = {}
+            multimodal_data["multimodal_embedding"] = mm_embeds
+            return token_ids.tolist(), {"multimodal_data": multimodal_data}
         else:
             return processed["input_ids"].squeeze().tolist(), {}
 
@@ -882,8 +885,12 @@ def forward(
         spec_metadata: Optional[SpecMetadata] = None,
         **kwargs,
     ) -> torch.Tensor:
-        mm_embed = kwargs.get("multi_modal_data", [])
-        if mm_embed:
+        multimodal_params = kwargs.get("multimodal_params", [])
+        if multimodal_params:
+            mm_embed = [
+                multimodal_param.multimodal_data["multimodal_embedding"]
+                for multimodal_param in multimodal_params
+            ]
             _, inputs_embeds = fuse_input_embeds(self.model.embed_tokens,
                                                  input_ids, mm_embed)
         return super().forward(attn_metadata,
diff --git a/tensorrt_llm/_torch/models/modeling_llava_next.py b/tensorrt_llm/_torch/models/modeling_llava_next.py
index 15f7f5452d3e..851b350b3638 100644
--- a/tensorrt_llm/_torch/models/modeling_llava_next.py
+++ b/tensorrt_llm/_torch/models/modeling_llava_next.py
@@ -210,8 +210,10 @@ def __call__(
         mm_features = torch.stack(
             [self._process(tensor) for tensor in mm_tensor])
         fused_input_ids, mm_features = self._postprocess(input_ids, mm_features)
+        multimodal_data = {}
+        multimodal_data["multimodal_embedding"] = mm_features
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mm_embedding": mm_features
+            "multimodal_data": multimodal_data
         }
 
 
@@ -273,7 +275,7 @@ def forward(
 
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.multimodal_embedding
+            multimodal_param.multimodal_data["multimodal_embedding"]
             for multimodal_param in multimodal_params
         ]
         assert mm_embed == [] or len(
diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 917752d12fad..5e4970d6d812 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -1,4 +1,5 @@
 import copy
+import os
 from typing import Any, Dict, List, Optional, Tuple, Union
 
 import torch
@@ -20,6 +21,8 @@
 from .modeling_multimodal_utils import fuse_input_embeds
 from .modeling_utils import register_auto_model
 
+DISAGG = os.getenv('TLLM_MULTIMODAL_DISAGGREGATED', '0') == '1'
+
 
 class Qwen2VLInputProcessorBase(InputProcessor):
 
@@ -322,7 +325,8 @@ def get_mrope_config(
         concat_cos_sin = concat_cos_sin.reshape(concat_cos_sin.shape[0], -1)
         mrope_config = {}
         mrope_config['mrope_rotary_cos_sin'] = concat_cos_sin.to('cpu')
-        mrope_config['mrope_position_deltas'] = mrope_position_deltas.to('cpu')
+        mrope_config['mrope_position_deltas'] = mrope_position_deltas.to(
+            'cpu').to(torch.int32)
         return mrope_config
 
     @torch.inference_mode()
@@ -364,11 +368,11 @@ def __call__(
             processed_inputs.get('video_grid_thw', None),
             processed_inputs.get('attention_mask', None),
             processed_inputs.get('second_per_grid_ts', None))
+        multimodal_data["mrope_config"] = mrope_config
 
         fused_input_ids = self._postprocess(input_ids[0])
 
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mrope_config": mrope_config,
             "multimodal_data": multimodal_data,
         }
 
@@ -411,16 +415,14 @@ def _parse_and_batch_multimodal_data(
 
         for multimodal_param in multimodal_params:
             # Process images if present
-            if "image" in multimodal_param.multimodal_data and multimodal_param.multimodal_data[
-                    "image"]:
+            if multimodal_param.multimodal_data.get("image") is not None:
                 pixel_values_list.append(
                     multimodal_param.multimodal_data["image"]["pixel_values"])
                 image_grid_thw_list.append(
                     multimodal_param.multimodal_data["image"]["image_grid_thw"])
 
             # Process videos if present
-            if "video" in multimodal_param.multimodal_data and multimodal_param.multimodal_data[
-                    "video"]:
+            if multimodal_param.multimodal_data.get("video") is not None:
                 pixel_values_videos_list.append(
                     multimodal_param.multimodal_data["video"]
                     ["pixel_values_videos"])
@@ -457,6 +459,8 @@ def forward(self, multimodal_params: List[MultimodalParams]):
 
         mm_content_data, mm_extra_data = self._parse_and_batch_multimodal_data(
             multimodal_params)
+        print(f"mm_content_data: {mm_content_data}")
+        print(f"mm_extra_data: {mm_extra_data}")
         pixel_values = mm_content_data.get("pixel_values", None)
         pixel_values_videos = mm_content_data.get("pixel_values_videos", None)
 
@@ -478,7 +482,6 @@ def forward(self, multimodal_params: List[MultimodalParams]):
             pixel_values_videos = pixel_values_videos.to(self.visual.dtype)
             embeds.append(
                 self.visual(pixel_values_videos, grid_thw=video_grid_thw))
-
         return embeds
 
 
@@ -526,16 +529,19 @@ def _parse_mrope_config(
         mrope_config = {}
         mrope_rotary_cos_sin_list = []
         mrope_position_deltas_list = []
-
         for multimodal_param in multimodal_params:
-            if hasattr(multimodal_param,
-                       'mrope_config') and multimodal_param.mrope_config:
-                if 'mrope_rotary_cos_sin' in multimodal_param.mrope_config:
+            if multimodal_param.multimodal_data and multimodal_param.multimodal_data.get(
+                    'mrope_config'):
+                if multimodal_param.multimodal_data['mrope_config'].get(
+                        'mrope_rotary_cos_sin') is not None:
                     mrope_rotary_cos_sin_list.append(
-                        multimodal_param.mrope_config['mrope_rotary_cos_sin'])
-                if 'mrope_position_deltas' in multimodal_param.mrope_config:
+                        multimodal_param.multimodal_data['mrope_config']
+                        ['mrope_rotary_cos_sin'])
+                if multimodal_param.multimodal_data['mrope_config'].get(
+                        'mrope_position_deltas') is not None:
                     mrope_position_deltas_list.append(
-                        multimodal_param.mrope_config['mrope_position_deltas'])
+                        multimodal_param.multimodal_data['mrope_config']
+                        ['mrope_position_deltas'])
 
         if mrope_rotary_cos_sin_list:
             mrope_config['mrope_rotary_cos_sin'] = torch.cat(
@@ -544,6 +550,8 @@ def _parse_mrope_config(
         if mrope_position_deltas_list:
             mrope_config['mrope_position_deltas'] = torch.cat(
                 mrope_position_deltas_list, dim=0)
+        print(f"mrope_config: {mrope_config}")
+        return mrope_config
 
     @torch.inference_mode()
     def forward(
@@ -568,8 +576,14 @@ def forward(
         mrope_config = {}
 
         if len(multimodal_params) > 0:
-            mm_embeds = self.mm_encoder.forward(
-                multimodal_params[:num_context_requests])
+            if not DISAGG:
+                mm_embeds = self.mm_encoder.forward(
+                    multimodal_params[:num_context_requests])
+            else:
+                mm_embeds = [
+                    multimodal_param.multimodal_data["multimodal_embedding"]
+                    for multimodal_param in multimodal_params
+                ]
             mrope_config = self._parse_mrope_config(multimodal_params)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
@@ -592,8 +606,9 @@ class Qwen2VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
-        self.mm_encoder = Qwen2VisionModelBase(model_config,
-                                               Qwen2VLForConditionalGeneration)
+        if not DISAGG:
+            self.mm_encoder = Qwen2VisionModelBase(
+                model_config, Qwen2VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
 
 
@@ -603,6 +618,7 @@ class Qwen2_5_VLModel(Qwen2VLModelBase):
 
     def __init__(self, model_config: ModelConfig[PretrainedConfig], *args,
                  **kwargs):
-        self.mm_encoder = Qwen2VisionModelBase(
-            model_config, Qwen2_5_VLForConditionalGeneration)
+        if not DISAGG:
+            self.mm_encoder = Qwen2VisionModelBase(
+                model_config, Qwen2_5_VLForConditionalGeneration)
         super().__init__(model_config, *args, **kwargs)
diff --git a/tensorrt_llm/_torch/models/modeling_vila.py b/tensorrt_llm/_torch/models/modeling_vila.py
index 3653f9c43fec..c27a88abf5f0 100644
--- a/tensorrt_llm/_torch/models/modeling_vila.py
+++ b/tensorrt_llm/_torch/models/modeling_vila.py
@@ -1107,8 +1107,10 @@ def __call__(
         )  # use_fast uses Pytorch GPU preprocessing, otherwise uses PIL CPU preprocessing
         mm_features = self._process(mm_tensor, block_sizes)
         fused_input_ids, mm_features = self._postprocess(input_ids, mm_features)
+        multimodal_data = {}
+        multimodal_data["multimodal_embedding"] = mm_features
         return fused_input_ids.to(torch.int32).tolist(), {
-            "mm_embedding": mm_features
+            "multimodal_data": multimodal_data
         }
 
 
@@ -1163,7 +1165,7 @@ def forward(
         num_context_requests, num_generation_requests = attn_metadata.num_contexts, attn_metadata.num_generations
         multimodal_params = kwargs.get("multimodal_params", [])
         mm_embed = [
-            multimodal_param.multimodal_embedding
+            multimodal_param.multimodal_data["multimodal_embedding"]
             for multimodal_param in multimodal_params
         ]
 
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 1eb49b64e1fc..adba493c8c4d 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1197,30 +1197,15 @@ def _prepare_tp_inputs(
             prompt_lengths.append(len(prompt_tokens))
             past_seen_token_num = begin_compute
             num_cached_tokens_per_seq.append(past_seen_token_num)
+            request.py_batch_idx = py_batch_idx(request)
 
-            multimodal_embedding = request.multimodal_embedding
-            if multimodal_embedding is not None:
-                # TODO: Visit later once we have the SharedTensor.
-                multimodal_embedding = multimodal_embedding.pin_memory(
-                ) if multimodal_embedding.device == 'cpu' else multimodal_embedding
-                multimodal_embedding = multimodal_embedding.to(
-                    'cuda', non_blocking=True)
-
-            mrope_rotary_cos_sin = request.mrope_rotary_cos_sin
-            if mrope_rotary_cos_sin is not None:
-                # TODO: Visit later once we have the SharedTensor.
-                mrope_rotary_cos_sin = mrope_rotary_cos_sin.pin_memory(
-                ) if mrope_rotary_cos_sin.device == 'cpu' else mrope_rotary_cos_sin
-                mrope_rotary_cos_sin = mrope_rotary_cos_sin.to(
-                    'cuda', non_blocking=True)
-
-            # Create MultimodalParams from request data
+            # Multimodal
             multimodal_params = MultimodalParams(
-                multimodal_embedding=multimodal_embedding,
-                mrope_config={'mrope_rotary_cos_sin': mrope_rotary_cos_sin}
-                if mrope_rotary_cos_sin is not None else {},
-                multimodal_data=request.py_multimodal_data,
-            )
+                multimodal_data=request.py_multimodal_data, )
+            multimodal_params.strip_for_context()
+            multimodal_params.to_device("multimodal_data",
+                                        "cuda",
+                                        pin_memory=True)
 
             if multimodal_params.has_content():
                 multimodal_params_list.append(multimodal_params)
@@ -1252,20 +1237,15 @@ def _prepare_tp_inputs(
                     extend_requests.append(request)
             else:
                 generation_requests.append(request)
-
-            # Handle generation request multimodal params
-            mrope_position_deltas = request.mrope_position_deltas
-            if mrope_position_deltas is not None:
-                mrope_position_deltas_tensor = torch.tensor(
-                    [mrope_position_deltas], dtype=torch.int32, pin_memory=True)
-                multimodal_params = MultimodalParams(
-                    mrope_config={
-                        'mrope_position_deltas':
-                        mrope_position_deltas_tensor.to('cuda',
-                                                        non_blocking=True)
-                    })
-                if multimodal_params.has_content():
-                    multimodal_params_list.append(multimodal_params)
+            # Multimodal
+            multimodal_params = MultimodalParams(
+                multimodal_data=request.py_multimodal_data, )
+            multimodal_params.strip_for_generation()
+            multimodal_params.to_device("multimodal_data",
+                                        "cuda",
+                                        pin_memory=True)
+            if multimodal_params.has_content():
+                multimodal_params_list.append(multimodal_params)
         extend_requests += extend_dummy_requests
 
         if not self._disable_overlap_scheduler and self.is_spec_decode:
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index 228410beae8a..38efdfefcd23 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -386,20 +386,12 @@ def _enqueue_request(self, request: GenerationRequest) -> int:
             prompt_token_ids = list(range(
                 vocab_size, vocab_size + pa_length)) + prompt_token_ids
 
-        # Multimodal related fields - simplified handling
+        # MULTIMODAL
+        # NOTE: Since, we only support PyTorch backend for multimodal, we will send multimodal_data through the 'py_multimodal_data' field
+        # except `multimodal_input` as it needs to go through the C++ runtime.
+        multimodal_input = None
         if request.multimodal_params is not None and request.multimodal_params.has_content(
         ):
-            # Create mrope_config if needed
-            mrope_config = None
-            if request.multimodal_params.mrope_config:
-                mrope_config = tllm.MropeConfig(
-                    mrope_rotary_cos_sin=request.multimodal_params.mrope_config.
-                    get('mrope_rotary_cos_sin'),
-                    mrope_position_deltas=request.multimodal_params.
-                    mrope_config.get('mrope_position_deltas'))
-
-            # Create multimodal_input for C++ if needed
-            multimodal_input = None
             if request.multimodal_params.multimodal_input is not None:
                 multimodal_input = tllm.MultimodalInput(
                     multimodal_hashes=request.multimodal_params.
@@ -408,13 +400,8 @@ def _enqueue_request(self, request: GenerationRequest) -> int:
                     multimodal_input.multimodal_positions,
                     multimodal_lengths=request.multimodal_params.
                     multimodal_input.multimodal_lengths)
-            multimodal_embedding = None
-            if request.multimodal_params.multimodal_embedding is not None:
-                multimodal_embedding = request.multimodal_params.multimodal_embedding
-        else:
-            multimodal_embedding = None
-            mrope_config = None
-            multimodal_input = None
+            # NOTE: Setting to None here to avoid sending multimodal_input again through the 'py_multimodal_data' field
+            request.multimodal_params.multimodal_input = None
 
         context_phase_params = None
         request_type = tllm.RequestType.REQUEST_TYPE_CONTEXT_AND_GENERATION
@@ -482,8 +469,8 @@ def _deduce_max_tokens(request: GenerationRequest,
                 lora_config=lora_config,
                 prompt_tuning_config=prompt_tuning_config,
                 multimodal_input=multimodal_input,
-                multimodal_embedding=multimodal_embedding,
-                mrope_config=mrope_config,
+                multimodal_embedding=None,
+                mrope_config=None,
                 logits_post_processor_name=(
                     tllm.Request.BATCHED_POST_PROCESSOR_NAME
                     if request.sampling_params.apply_batched_logits_processor
@@ -494,9 +481,8 @@ def _deduce_max_tokens(request: GenerationRequest,
                 context_phase_params=context_phase_params,
                 type=request_type)
 
-            if self._is_pytorch_backend:
-                # For PyTorch backend, attach the raw multimodal data
-                if request.multimodal_params is not None and request.multimodal_params.multimodal_data:
+            if self._is_pytorch_backend and request.multimodal_params is not None:
+                if request.multimodal_params.multimodal_data is not None:
                     executor_request.py_multimodal_data = request.multimodal_params.multimodal_data
 
             if self._is_pytorch_backend and request.sampling_params.logits_processor:
diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index 8d8b7962597a..ca8dfd552b5b 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -90,21 +90,26 @@ class MultimodalParams:
     providing a clean interface for handling multimodal inputs across different models.
     """
 
-    # Core multimodal data
     multimodal_input: Optional[MultimodalInput] = None
-    """Multimodal input data with hashing information for caching and deduplication."""
+    """Multimodal input data with hashing information."""
 
-    multimodal_embedding: Optional[torch.Tensor] = None
-    """Pre-computed multimodal embeddings from vision encoder."""
+    multimodal_data: Optional[Dict[str, Any]] = field(default_factory=dict)
+    """Processed multimodal data after AutoProcessor's process() by modality, multimodal_embedding and mrope_config.
+    It should be in the form of {"mrope_config": {"mrope_rotary_cos_sin": torch.Tensor, "mrope_position_deltas": torch.Tensor},
+                                 "multimodal_embedding": torch.Tensor,
+                                 "modality": {item_str: item_data}}
+
+    mrope_config: Optional[Dict[str, Any]]: Multimodal rotary position embedding config (used by Qwen2/2.5-VL).
+    multimodal_embedding: Optional[torch.Tensor]: Pre-computed multimodal embeddings from vision encoder.
+    modality: Optional[Dict[str, Any]]: Multimodal data by modality. For example, image, video, etc.
 
-    multimodal_data: Optional[Dict[str, Dict[str,
-                                             Union[torch.Tensor,
-                                                   List[Any]]]]] = field(
-                                                       default_factory=dict)
-    """Processed multimodal data after AutoProcessor's process() by modality (e.g., image pixels, video pixel values).
-    It should be in the form of {modality: {item_str: item_data}}
     e.g.
     {
+        "mrope_config": {
+            "mrope_rotary_cos_sin": torch.Tensor(),
+            "mrope_position_deltas": torch.Tensor(),
+        },
+        "multimodal_embedding": torch.Tensor(),
         "image": {
             "pixel_values": torch.Tensor(),
             "image_height": torch.Tensor() or List[int],
@@ -115,25 +120,74 @@ class MultimodalParams:
             "video_height": torch.Tensor() or List[int],
             "video_width": torch.Tensor() or List[int]
         },
+        ...
     }
     """
 
-    # Model-specific configurations
-    mrope_config: Optional[Dict[str, Any]] = None
-    """Multimodal rotary position embedding config (used by Qwen2-VL)."""
-
     def __post_init__(self):
         """Ensure default values are properly set."""
         if self.multimodal_data is None:
             self.multimodal_data = {}
-        if self.mrope_config is None:
-            self.mrope_config = {}
+
+    def to_device(self, element: str, device: str, pin_memory: bool = False):
+
+        def _to_device(
+            input_tensor: Union[torch.Tensor, List, dict, None],
+            pin_memory: bool = False,
+        ) -> Union[torch.Tensor, List, dict, None]:
+            if input_tensor is None:
+                return None
+            elif isinstance(input_tensor, list):
+                return [_to_device(item) for item in input_tensor]
+            elif isinstance(input_tensor, dict):
+                return {
+                    key: _to_device(value)
+                    for key, value in input_tensor.items()
+                }
+            elif isinstance(input_tensor, torch.Tensor):
+                if pin_memory:
+                    return input_tensor.pin_memory().to(device)
+                else:
+                    return input_tensor.to(device)
+
+        if element == "multimodal_data":
+            self.multimodal_data = _to_device(self.multimodal_data)
+        elif element == "multimodal_input":
+            self.multimodal_input = _to_device(self.multimodal_input)
+        else:
+            print("MultimodalParams: Unsupported element to move to device: ",
+                  element)
+
+    def strip_for_context(self):
+        """Strip multimodal data for context mode - remove only mrope_position_deltas."""
+        if self.multimodal_data and 'mrope_config' in self.multimodal_data:
+            mrope_config = self.multimodal_data['mrope_config']
+            if 'mrope_position_deltas' in mrope_config:
+                # Remove only mrope_position_deltas, keep everything else
+                del mrope_config['mrope_position_deltas']
+
+    def strip_for_generation(self):
+        """Strip multimodal data for generation mode - keep only mrope_position_deltas."""
+        if self.multimodal_data:
+            # In generation mode, only keep mrope_config['mrope_position_deltas']
+            # and erase everything else
+            mrope_position_deltas = None
+            if 'mrope_config' in self.multimodal_data:
+                mrope_config = self.multimodal_data['mrope_config']
+                if 'mrope_position_deltas' in mrope_config:
+                    mrope_position_deltas = mrope_config[
+                        'mrope_position_deltas']
+
+            # Clear all multimodal_data and only keep mrope_position_deltas if it exists
+            self.multimodal_data = {}
+            if mrope_position_deltas is not None:
+                self.multimodal_data['mrope_config'] = {
+                    'mrope_position_deltas': mrope_position_deltas
+                }
 
     def has_content(self) -> bool:
         """Check if this object contains any multimodal data."""
-        return bool(self.multimodal_input
-                    or self.multimodal_embedding is not None
-                    or self.multimodal_data or self.mrope_config)
+        return bool(self.multimodal_input or self.multimodal_data)
 
 
 # adopt from vllm : https://github.com/vllm-project/vllm/blob/main/vllm/vllm/multimodal/hash.py
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index 86d1be4b8882..6f3adcbda2f6 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -386,9 +386,6 @@ def generate_async(
                 query_token_ids = extra_processed_inputs.get('query_token_ids')
                 # Create unified MultimodalParams
                 multimodal_params = MultimodalParams(
-                    multimodal_embedding=extra_processed_inputs.get(
-                        'mm_embedding'),
-                    mrope_config=extra_processed_inputs.get('mrope_config'),
                     multimodal_input=extra_processed_inputs.get(
                         'multimodal_input'),
                     multimodal_data=extra_processed_inputs.get(

From 00f82a070168711184e969cfa1ffe32d88b81f36 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 4 Jul 2025 15:31:07 +0900
Subject: [PATCH 09/14] add more description to MultimodalParams

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 tensorrt_llm/inputs/multimodal.py | 158 +++++++++++++++++-------------
 1 file changed, 91 insertions(+), 67 deletions(-)

diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index ca8dfd552b5b..a6b29a9f0183 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -88,48 +88,52 @@ class MultimodalParams:
 
     This class encapsulates all multimodal-related data that flows through the system,
     providing a clean interface for handling multimodal inputs across different models.
+
+    Attributes:
+        multimodal_input: Multimodal input data with hashing information.
+        multimodal_data: Processed multimodal data containing embeddings, configurations,
+                        and modality-specific data organized by type.
+
+    Structure of multimodal_data:
+        {
+            "mrope_config": {
+                "mrope_rotary_cos_sin": torch.Tensor,    # Rotary embeddings (Qwen2/2.5-VL)
+                "mrope_position_deltas": torch.Tensor,   # Position deltas (Qwen2/2.5-VL)
+            },
+            "multimodal_embedding": torch.Tensor,        # Pre-computed vision embeddings
+            "image": {
+                "pixel_values": torch.Tensor,
+                "image_height": torch.Tensor | List[int],
+                "image_width": torch.Tensor | List[int],
+            },
+            "video": {
+                "pixel_values": torch.Tensor,
+                "video_height": torch.Tensor | List[int],
+                "video_width": torch.Tensor | List[int],
+            },
+            # ... other modalities
+        }
     """
 
     multimodal_input: Optional[MultimodalInput] = None
-    """Multimodal input data with hashing information."""
-
     multimodal_data: Optional[Dict[str, Any]] = field(default_factory=dict)
-    """Processed multimodal data after AutoProcessor's process() by modality, multimodal_embedding and mrope_config.
-    It should be in the form of {"mrope_config": {"mrope_rotary_cos_sin": torch.Tensor, "mrope_position_deltas": torch.Tensor},
-                                 "multimodal_embedding": torch.Tensor,
-                                 "modality": {item_str: item_data}}
-
-    mrope_config: Optional[Dict[str, Any]]: Multimodal rotary position embedding config (used by Qwen2/2.5-VL).
-    multimodal_embedding: Optional[torch.Tensor]: Pre-computed multimodal embeddings from vision encoder.
-    modality: Optional[Dict[str, Any]]: Multimodal data by modality. For example, image, video, etc.
-
-    e.g.
-    {
-        "mrope_config": {
-            "mrope_rotary_cos_sin": torch.Tensor(),
-            "mrope_position_deltas": torch.Tensor(),
-        },
-        "multimodal_embedding": torch.Tensor(),
-        "image": {
-            "pixel_values": torch.Tensor(),
-            "image_height": torch.Tensor() or List[int],
-            "image_width": torch.Tensor() or List[int]
-        },
-        "video": {
-            "pixel_values": torch.Tensor(),
-            "video_height": torch.Tensor() or List[int],
-            "video_width": torch.Tensor() or List[int]
-        },
-        ...
-    }
-    """
 
     def __post_init__(self):
         """Ensure default values are properly set."""
         if self.multimodal_data is None:
             self.multimodal_data = {}
 
-    def to_device(self, element: str, device: str, pin_memory: bool = False):
+    def to_device(self,
+                  element: str,
+                  device: str,
+                  pin_memory: bool = False) -> None:
+        """Move specified multimodal data element to target device.
+
+        Args:
+            element: Element to move ("multimodal_data" or "multimodal_input")
+            device: Target device (e.g., "cuda", "cpu")
+            pin_memory: Whether to pin memory for faster transfers
+        """
 
         def _to_device(
             input_tensor: Union[torch.Tensor, List, dict, None],
@@ -138,52 +142,72 @@ def _to_device(
             if input_tensor is None:
                 return None
             elif isinstance(input_tensor, list):
-                return [_to_device(item) for item in input_tensor]
+                return [_to_device(item, pin_memory) for item in input_tensor]
             elif isinstance(input_tensor, dict):
                 return {
-                    key: _to_device(value)
+                    key: _to_device(value, pin_memory)
                     for key, value in input_tensor.items()
                 }
             elif isinstance(input_tensor, torch.Tensor):
-                if pin_memory:
-                    return input_tensor.pin_memory().to(device)
+                if pin_memory and input_tensor.device.type == 'cpu':
+                    return input_tensor.pin_memory().to(device,
+                                                        non_blocking=True)
                 else:
-                    return input_tensor.to(device)
+                    return input_tensor.to(device, non_blocking=True)
+            else:
+                return input_tensor
 
         if element == "multimodal_data":
-            self.multimodal_data = _to_device(self.multimodal_data)
+            self.multimodal_data = _to_device(self.multimodal_data, pin_memory)
         elif element == "multimodal_input":
-            self.multimodal_input = _to_device(self.multimodal_input)
+            self.multimodal_input = _to_device(self.multimodal_input,
+                                               pin_memory)
         else:
-            print("MultimodalParams: Unsupported element to move to device: ",
-                  element)
-
-    def strip_for_context(self):
-        """Strip multimodal data for context mode - remove only mrope_position_deltas."""
-        if self.multimodal_data and 'mrope_config' in self.multimodal_data:
+            print(
+                f"MultimodalParams: Unsupported element '{element}' to move to device. "
+                f"Supported elements: 'multimodal_data', 'multimodal_input'")
+
+    def strip_for_context(self) -> None:
+        """Strip multimodal data for context processing.
+
+        Removes only mrope_position_deltas while keeping all other multimodal data
+        (embeddings, images, etc.) needed for context phase processing.
+        """
+        if not (self.multimodal_data
+                and 'mrope_config' in self.multimodal_data):
+            return
+
+        mrope_config = self.multimodal_data['mrope_config']
+        if 'mrope_position_deltas' in mrope_config:
+            del mrope_config['mrope_position_deltas']
+
+            # Clean up empty mrope_config
+            if not mrope_config:
+                del self.multimodal_data['mrope_config']
+
+    def strip_for_generation(self) -> None:
+        """Strip multimodal data for generation processing.
+
+        Keeps only mrope_position_deltas and removes all other multimodal data
+        (embeddings, images, etc.) as they're not needed during generation.
+        """
+        if not self.multimodal_data:
+            return
+
+        # Extract mrope_position_deltas before clearing
+        mrope_position_deltas = None
+        if 'mrope_config' in self.multimodal_data:
             mrope_config = self.multimodal_data['mrope_config']
-            if 'mrope_position_deltas' in mrope_config:
-                # Remove only mrope_position_deltas, keep everything else
-                del mrope_config['mrope_position_deltas']
-
-    def strip_for_generation(self):
-        """Strip multimodal data for generation mode - keep only mrope_position_deltas."""
-        if self.multimodal_data:
-            # In generation mode, only keep mrope_config['mrope_position_deltas']
-            # and erase everything else
-            mrope_position_deltas = None
-            if 'mrope_config' in self.multimodal_data:
-                mrope_config = self.multimodal_data['mrope_config']
-                if 'mrope_position_deltas' in mrope_config:
-                    mrope_position_deltas = mrope_config[
-                        'mrope_position_deltas']
-
-            # Clear all multimodal_data and only keep mrope_position_deltas if it exists
-            self.multimodal_data = {}
-            if mrope_position_deltas is not None:
-                self.multimodal_data['mrope_config'] = {
-                    'mrope_position_deltas': mrope_position_deltas
-                }
+            if isinstance(mrope_config,
+                          dict) and 'mrope_position_deltas' in mrope_config:
+                mrope_position_deltas = mrope_config['mrope_position_deltas']
+
+        # Clear all data and restore only position deltas if they exist
+        self.multimodal_data = {}
+        if mrope_position_deltas is not None:
+            self.multimodal_data['mrope_config'] = {
+                'mrope_position_deltas': mrope_position_deltas
+            }
 
     def has_content(self) -> bool:
         """Check if this object contains any multimodal data."""

From a7a4e19458e9ec0fc866384737de0375b1a95111 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 4 Jul 2025 16:55:45 +0900
Subject: [PATCH 10/14] rebase fix

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/model_engine.py | 1 -
 1 file changed, 1 deletion(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index adba493c8c4d..2dbbe5af7771 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1197,7 +1197,6 @@ def _prepare_tp_inputs(
             prompt_lengths.append(len(prompt_tokens))
             past_seen_token_num = begin_compute
             num_cached_tokens_per_seq.append(past_seen_token_num)
-            request.py_batch_idx = py_batch_idx(request)
 
             # Multimodal
             multimodal_params = MultimodalParams(

From 777a83df443f2ef706af85fe6ac8ddcea26a5236 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 4 Jul 2025 19:18:41 +0900
Subject: [PATCH 11/14] mrope_position_deltas fix

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/models/modeling_qwen2vl.py         | 76 +++++++++++--------
 .../_torch/pyexecutor/model_engine.py         |  5 +-
 2 files changed, 46 insertions(+), 35 deletions(-)

diff --git a/tensorrt_llm/_torch/models/modeling_qwen2vl.py b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
index 5e4970d6d812..1f9bfb803253 100644
--- a/tensorrt_llm/_torch/models/modeling_qwen2vl.py
+++ b/tensorrt_llm/_torch/models/modeling_qwen2vl.py
@@ -459,8 +459,6 @@ def forward(self, multimodal_params: List[MultimodalParams]):
 
         mm_content_data, mm_extra_data = self._parse_and_batch_multimodal_data(
             multimodal_params)
-        print(f"mm_content_data: {mm_content_data}")
-        print(f"mm_extra_data: {mm_extra_data}")
         pixel_values = mm_content_data.get("pixel_values", None)
         pixel_values_videos = mm_content_data.get("pixel_values_videos", None)
 
@@ -523,35 +521,47 @@ def post_config(self):
         self.config = self.llm.config
         self.model_config.pretrained_config = self.llm.config
 
-    def _parse_mrope_config(
-            self, multimodal_params: List[MultimodalParams]
-    ) -> dict[str, torch.Tensor]:
-        mrope_config = {}
-        mrope_rotary_cos_sin_list = []
-        mrope_position_deltas_list = []
-        for multimodal_param in multimodal_params:
-            if multimodal_param.multimodal_data and multimodal_param.multimodal_data.get(
-                    'mrope_config'):
-                if multimodal_param.multimodal_data['mrope_config'].get(
-                        'mrope_rotary_cos_sin') is not None:
-                    mrope_rotary_cos_sin_list.append(
-                        multimodal_param.multimodal_data['mrope_config']
-                        ['mrope_rotary_cos_sin'])
-                if multimodal_param.multimodal_data['mrope_config'].get(
-                        'mrope_position_deltas') is not None:
-                    mrope_position_deltas_list.append(
-                        multimodal_param.multimodal_data['mrope_config']
-                        ['mrope_position_deltas'])
-
-        if mrope_rotary_cos_sin_list:
-            mrope_config['mrope_rotary_cos_sin'] = torch.cat(
-                mrope_rotary_cos_sin_list, dim=0)
-
-        if mrope_position_deltas_list:
-            mrope_config['mrope_position_deltas'] = torch.cat(
-                mrope_position_deltas_list, dim=0)
-        print(f"mrope_config: {mrope_config}")
-        return mrope_config
+    def _parse_and_concat_mrope_config(
+            self, multimodal_params: List[MultimodalParams],
+            num_context_requests: int,
+            num_generation_requests: int) -> dict[str, torch.Tensor]:
+        """
+        Parse and concatenate mrope configuration from multimodal parameters.
+        """
+
+        mrope_configs = [
+            param.multimodal_data.get('mrope_config')
+            for param in multimodal_params if param.multimodal_data
+            and param.multimodal_data.get('mrope_config')
+        ]
+        if not mrope_configs:
+            return {}
+
+        batched_mrope_config = {}
+        if num_context_requests > 0:
+            cos_sin_tensors = [
+                config['mrope_rotary_cos_sin']
+                for config in mrope_configs[:num_context_requests]
+                if config.get('mrope_rotary_cos_sin') is not None
+            ]
+            if cos_sin_tensors:
+                batched_mrope_config['mrope_rotary_cos_sin'] = torch.cat(
+                    cos_sin_tensors, dim=0)
+
+        if num_generation_requests > 0:
+            generation_mrope_configs = mrope_configs[
+                -num_generation_requests:] if len(
+                    mrope_configs) >= num_generation_requests else mrope_configs
+            position_delta_tensors = [
+                config['mrope_position_deltas']
+                for config in generation_mrope_configs
+                if config.get('mrope_position_deltas') is not None
+            ]
+            if position_delta_tensors:
+                batched_mrope_config['mrope_position_deltas'] = torch.cat(
+                    position_delta_tensors, dim=0)
+
+        return batched_mrope_config
 
     @torch.inference_mode()
     def forward(
@@ -584,7 +594,9 @@ def forward(
                     multimodal_param.multimodal_data["multimodal_embedding"]
                     for multimodal_param in multimodal_params
                 ]
-            mrope_config = self._parse_mrope_config(multimodal_params)
+            mrope_config = self._parse_and_concat_mrope_config(
+                multimodal_params, num_context_requests,
+                num_generation_requests)
 
         input_ids, input_embeds = fuse_input_embeds(self.llm.model.embed_tokens,
                                                     input_ids, mm_embeds)
diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 2dbbe5af7771..bf112b6ff56f 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1200,8 +1200,7 @@ def _prepare_tp_inputs(
 
             # Multimodal
             multimodal_params = MultimodalParams(
-                multimodal_data=request.py_multimodal_data, )
-            multimodal_params.strip_for_context()
+                multimodal_data=request.py_multimodal_data)
             multimodal_params.to_device("multimodal_data",
                                         "cuda",
                                         pin_memory=True)
@@ -1238,7 +1237,7 @@ def _prepare_tp_inputs(
                 generation_requests.append(request)
             # Multimodal
             multimodal_params = MultimodalParams(
-                multimodal_data=request.py_multimodal_data, )
+                multimodal_data=request.py_multimodal_data)
             multimodal_params.strip_for_generation()
             multimodal_params.to_device("multimodal_data",
                                         "cuda",

From 82cf62ebd693b525e0d351c46c72a3ea79098f71 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Sat, 5 Jul 2025 00:03:32 +0900
Subject: [PATCH 12/14] change image_url to video_url on video test

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py | 2 +-
 1 file changed, 1 insertion(+), 1 deletion(-)

diff --git a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
index 66ff8d59e176..171d6d5bb3fc 100644
--- a/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
+++ b/tests/unittest/llmapi/apps/_test_openai_chat_multimodal.py
@@ -188,7 +188,7 @@ def test_single_chat_session_video(client: openai.OpenAI, model_name: str):
             "type": "text",
             "text": content_text
         }, {
-            "type": "image_url",
+            "type": "video_url",
             "video_url": {
                 "url": video_url
             }

From dcf6ab10a6c0c027fb7097af9f90f06743f66e2c Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Tue, 8 Jul 2025 01:37:16 +0900
Subject: [PATCH 13/14] rebase fix

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 tensorrt_llm/_torch/pyexecutor/model_engine.py | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index bf112b6ff56f..673ac0c6a2ca 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1176,10 +1176,9 @@ def _prepare_tp_inputs(
         gather_ids = []
         position_ids = []  # per sequence
         num_cached_tokens_per_seq = []  # per sequence
-        multi_modal_data = []
         draft_tokens = []
         draft_lens = []
-        mrope_config = defaultdict(list)
+        multimodal_params_list = []
         gen_request_seq_slots = []  # per generation request
 
         for request in scheduled_requests.context_requests:

From 8c50d0208647eaa020dc0ca26b229b609e057ccf Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Tue, 8 Jul 2025 01:56:49 +0900
Subject: [PATCH 14/14] add comment

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 tensorrt_llm/executor/worker.py | 1 +
 1 file changed, 1 insertion(+)

diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index 38efdfefcd23..da90fc8fe931 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -469,6 +469,7 @@ def _deduce_max_tokens(request: GenerationRequest,
                 lora_config=lora_config,
                 prompt_tuning_config=prompt_tuning_config,
                 multimodal_input=multimodal_input,
+                #NOTE: `multimodal_embedding` and `mrope_config` will be in MultimodalParams.multimodal_data. And this will be handled below by `py_multimodal_data`.
                 multimodal_embedding=None,
                 mrope_config=None,
                 logits_post_processor_name=(