From 2abd27d933d49b16176505c17a1e14b496eac3d4 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Tue, 22 Jul 2025 17:48:37 +0900
Subject: [PATCH 1/9] feat: support SharedTensor on MultimodalParams

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         |  18 +-
 tensorrt_llm/inputs/multimodal.py             | 376 +++++++++---------
 tensorrt_llm/llmapi/llm.py                    |   2 +
 3 files changed, 196 insertions(+), 200 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index e6044133f939..7fcf1924f1df 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1227,11 +1227,15 @@ def _prepare_tp_inputs(
             multimodal_params = MultimodalParams(
                 multimodal_data=request.py_multimodal_data,
                 multimodal_runtime=py_multimodal_runtime)
-            multimodal_params.to_device("multimodal_data",
-                                        "cuda",
-                                        pin_memory=True)
 
             if multimodal_params.has_content():
+                multimodal_params.from_shared_tensor("multimodal_data")
+                multimodal_params.to_device("multimodal_data",
+                                            "cuda",
+                                            pin_memory=True)
+                # re-assign the recovered SharedTensorContainer to the request for generation usage
+                setattr(request, "py_multimodal_data",
+                        multimodal_params.multimodal_data)
                 multimodal_params_list.append(multimodal_params)
 
             request.py_batch_idx = request.py_seq_slot
@@ -1264,11 +1268,11 @@ def _prepare_tp_inputs(
             # 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.strip_for_generation()
+                # re-assign the recovered SharedTensorContainer to the request after strip_for_generation
+                setattr(request, "py_multimodal_data",
+                        multimodal_params.multimodal_data)
                 multimodal_params_list.append(multimodal_params)
         extend_requests += extend_dummy_requests
 
diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index 936890658793..fd647c6fa53e 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -190,237 +190,227 @@ def __post_init__(self):
         if self.multimodal_data is None:
             self.multimodal_data = {}
 
-    def to_device(self,
-                  element: str,
-                  device: str,
-                  pin_memory: bool = False) -> None:
-        """Move specified multimodal data element to target device.
+    def _is_shared_tensor_dict(self, obj: Any) -> bool:
+        """Check if an object is a shared tensor dictionary.
 
         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
+            obj: Object to check
+
+        Returns:
+            True if the object is a shared tensor dictionary, False otherwise
         """
+        if not isinstance(obj, dict):
+            return False
+
+        # Check for required keys that uniquely identify a shared tensor dict
+        required_keys = {'method_key'}
+        if not required_keys.issubset(obj.keys()):
+            return False
+
+        # Additional validation based on method_key
+        method_key = obj.get('method_key')
+        if method_key == 1:  # CUDA tensor
+            cuda_keys = {'tensor_size', 'storage_handle', 'storage_device'}
+            return cuda_keys.issubset(obj.keys())
+        elif method_key == 2:  # CPU tensor
+            cpu_keys = {'tensor_size', 'storage_handle', 'manager_handle'}
+            return cpu_keys.issubset(obj.keys())
+
+        return False
+
+    def _apply_tensor_operation(
+            self, input_data: Union[torch.Tensor, List, dict, None],
+            operation: str, **kwargs) -> Union[torch.Tensor, List, dict, None]:
+        """Apply tensor operations recursively to nested data structures.
+
+        This method handles three types of operations:
+        - "to_shared_tensor": Convert tensors to shared tensor dictionaries
+        - "from_shared_tensor": Convert shared tensor dictionaries back to tensors
+        - "to_device": Move tensors to specified device
 
-        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, pin_memory) for item in input_tensor]
-            elif isinstance(input_tensor, dict):
+        Args:
+            input_data: Input data structure (tensor, list, dict, or None)
+            operation: Operation to apply
+            **kwargs: Additional arguments for the operation
+
+        Returns:
+            Transformed data structure
+        """
+        # Handle None case
+        if input_data is None:
+            return None
+
+        # Handle list case - recursively process each element
+        if isinstance(input_data, list):
+            return [
+                self._apply_tensor_operation(item, operation, **kwargs)
+                for item in input_data
+            ]
+
+        # Handle dictionary case
+        if isinstance(input_data, dict):
+            if operation == "from_shared_tensor" and self._is_shared_tensor_dict(
+                    input_data):
+                # Convert shared tensor dict back to tensor
+                try:
+                    from tensorrt_llm._torch.shared_tensor import \
+                        SharedTensorContainer
+
+                    # print(f"input_data: {input_data}")
+                    return SharedTensorContainer.from_dict(
+                        input_data).get_local_view()
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to restore tensor from shared tensor dict: {e}"
+                    )
+            else:
+                # Regular dictionary - recursively process values
                 return {
-                    key: _to_device(value, pin_memory)
-                    for key, value in input_tensor.items()
+                    key: self._apply_tensor_operation(value, operation,
+                                                      **kwargs)
+                    for key, value in input_data.items()
                 }
-            elif isinstance(input_tensor, torch.Tensor):
-                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, non_blocking=True)
-            else:
-                return input_tensor
 
-        if element == "multimodal_data":
-            self.multimodal_data = _to_device(self.multimodal_data, pin_memory)
-        elif element == "multimodal_input":
-            self.multimodal_input = _to_device(self.multimodal_input,
-                                               pin_memory)
-        else:
-            print(
-                f"MultimodalParams: Unsupported element '{element}' to move to device. "
-                f"Supported elements: 'multimodal_data', 'multimodal_input'")
+        # Handle tensor case
+        if isinstance(input_data, torch.Tensor):
+            if operation == "to_shared_tensor":
+                try:
+                    from tensorrt_llm._torch.shared_tensor import \
+                        SharedTensorContainer
+                    return SharedTensorContainer.from_tensor(
+                        input_data).dump_to_dict()
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to convert tensor to shared tensor: {e}")
+            elif operation == "to_device":
+                device = kwargs.get('device')
+                if device is None:
+                    raise ValueError(
+                        "Device must be specified for 'to_device' operation")
+
+                pin_memory = kwargs.get('pin_memory', False)
+                try:
+                    if pin_memory and input_data.device.type == 'cpu':
+                        return input_data.pin_memory().to(device,
+                                                          non_blocking=True)
+                    else:
+                        return input_data.to(device, non_blocking=True)
+                except Exception as e:
+                    raise RuntimeError(
+                        f"Failed to move tensor to device {device}: {e}")
+
+        # For any other type, return as-is
+        return input_data
+
+    def _validate_element(self, element: str,
+                          supported_elements: List[str]) -> None:
+        """Validate that the element is supported.
+
+        Args:
+            element: Element name to validate
+            supported_elements: List of supported element names
 
-    def to_handle(self, element: str, key: Optional[str] = None) -> None:
-        """Convert multimodal data to tensor handle.
+        Raises:
+            ValueError: If element is not supported
+        """
+        if element not in supported_elements:
+            raise ValueError(
+                f"Unsupported element '{element}'. "
+                f"Supported elements: {', '.join(repr(e) for e in supported_elements)}"
+            )
 
-        Converts torch.Tensor objects to SharedTensorContainer handles (serializable dictionaries)
-        for efficient IPC. This function is a in-place operation.
+    def to_shared_tensor(self, element: str) -> None:
+        """Move specified multimodal data element to shared tensor.
 
         Args:
-            element: Element to convert ("multimodal_data" or "multimodal_input")
-            key: Specific key to convert. If None, converts all tensor values in multimodal_data.
-                 Defaults to None.
-
-        Example:
-            # Convert all tensors in multimodal_data to handles
-            params.to_handle("multimodal_data", key=None)
+            element: Element to move ("multimodal_data" or "multimodal_input")
 
-            # Convert only multimodal_embedding section tensors to handles
-            params.to_handle("multimodal_data", key="multimodal_embedding")
+        Raises:
+            ValueError: If element is not supported
+            RuntimeError: If tensor conversion fails
         """
-        # Lazy import to avoid circular dependency
-        from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
-
-        def _to_tensor_handle(data):
-            for k, v in data.items():
-                if isinstance(v, torch.Tensor):
-                    # Convert tensor to handle
-                    handle = SharedTensorContainer.from_tensor(v).dump_to_dict()
-                    data[k] = handle
-                elif isinstance(v, dict):
-                    _to_tensor_handle(v)
-                elif isinstance(v, list):
-                    for i, item in enumerate(v):
-                        if isinstance(item, torch.Tensor):
-                            handle = SharedTensorContainer.from_tensor(
-                                item).dump_to_dict()
-                            v[i] = handle
-
-        if element == "multimodal_data":
-            if self.multimodal_data is None:
-                return
-            if key is None:
-                _to_tensor_handle(self.multimodal_data)
-            else:
-                if key not in self.multimodal_data:
-                    return  # no-op if key not found
-
-                value = self.multimodal_data[key]
-                if isinstance(value, torch.Tensor):
-                    handle = SharedTensorContainer.from_tensor(
-                        value).dump_to_dict()
-                    self.multimodal_data[key] = handle
-                elif isinstance(value, dict):
-                    _to_tensor_handle(value)
-                else:
-                    raise ValueError(
-                        f"Unsupported value type for multimodal_data: {type(value)}"
-                    )
-        elif element == "multimodal_input":
-            # No-op for multimodal_input
-            return
-        else:
-            raise ValueError(
-                f"Unsupported element '{element}' to convert to handle.")
+        supported_elements = ["multimodal_data", "multimodal_input"]
+        self._validate_element(element, supported_elements)
+
+        data = getattr(self, element)
+        if data is None:
+            return  # Nothing to convert
 
-    def to_tensor(self, element: str, key: Optional[str] = None) -> None:
-        """Convert multimodal tensor handles back to tensors. This is the dual operation to to_handle.
+        transformed_data = self._apply_tensor_operation(data,
+                                                        "to_shared_tensor")
+        setattr(self, element, transformed_data)
 
-        Converts SharedTensorContainer handles (serializable dictionaries) back to torch.Tensor objects
-        for local computation. This function performs in-place modifications to the multimodal_data.
+    def from_shared_tensor(self, element: str, mode: str = "context") -> None:
+        """Move specified multimodal data element from shared tensor.
 
         Args:
-            element: Element to convert ("multimodal_data" or "multimodal_input")
-            key: Specific key to convert. If None, converts all tensor handles in multimodal_data.
-                 Defaults to None.
+            element: Element to restore ("multimodal_data" or "multimodal_input")
+            mode: Recovery mode - "context" or "generation"
+                - "context": Recover all shared tensors except mrope_position_deltas
+                - "generation": Recover only mrope_position_deltas
+
+        Raises:
+            ValueError: If element or mode is not supported
+            RuntimeError: If tensor restoration fails
+        """
+        supported_elements = ["multimodal_data", "multimodal_input"]
 
-        Example:
-            # Convert all handles back to tensors
-            params.to_tensor("multimodal_data", key=None)
+        self._validate_element(element, supported_elements)
 
-            # Convert only multimodal_embedding section handles back to tensors
-            params.to_tensor("multimodal_data", key="multimodal_embedding")
-        """
-        # Lazy import to avoid circular dependency
-        from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
-
-        def _to_tensor(data):
-            for k, v in data.items():
-                if isinstance(v, dict) and 'method_key' in v:
-                    # This is a tensor handle (dict with method_key)
-                    try:
-                        tensor = SharedTensorContainer.from_dict(
-                            v).get_local_view()
-                        data[k] = tensor
-                    except Exception as e:
-                        raise ValueError(
-                            f"Failed to convert handle to tensor for key '{k}': {e}"
-                        )
-                elif isinstance(v, dict):
-                    _to_tensor(v)
-                elif isinstance(v, list):
-                    for i, item in enumerate(v):
-                        if isinstance(item, dict) and 'method_key' in item:
-                            try:
-                                tensor = SharedTensorContainer.from_dict(
-                                    item).get_local_view()
-                                v[i] = tensor
-                            except Exception as e:
-                                raise ValueError(
-                                    f"Failed to convert handle to tensor in list at index {i}: {e}"
-                                )
-
-        if element == "multimodal_data":
-            if self.multimodal_data is None:
-                return
-
-            if key is None:
-                _to_tensor(self.multimodal_data)
-            else:
-                if key not in self.multimodal_data:
-                    return  # no-op if key not found
-
-                value = self.multimodal_data[key]
-                if isinstance(
-                        value, dict
-                ) and 'method_key' in value:  # This is a tensor handle
-                    try:
-                        tensor = SharedTensorContainer.from_dict(
-                            value).get_local_view()
-                        self.multimodal_data[key] = tensor
-                    except Exception as e:
-                        raise ValueError(
-                            f"Failed to convert handle to tensor for key '{key}': {e}"
-                        )
-                elif isinstance(value, dict):
-                    _to_tensor(value)
-                else:
-                    raise ValueError(
-                        f"Unsupported value type for multimodal_data: {type(value)}"
-                    )
+        data = getattr(self, element)
+        if data is None:
+            return  # Nothing to restore
 
-        elif element == "multimodal_input":
-            # No-op for multimodal_input
-            return
-        else:
-            raise ValueError(
-                f"Unsupported element '{element}' to convert to tensor.")
+        restored_data = self._apply_tensor_operation(data, "from_shared_tensor")
+        setattr(self, element, restored_data)
+        # print(f"Restored {element} from shared tensor, restored_data: {restored_data}")
 
-    def strip_for_context(self) -> None:
-        """Strip multimodal data for context processing.
+    def to_device(self,
+                  element: str,
+                  device: str,
+                  pin_memory: bool = False) -> None:
+        """Move specified multimodal data element to target device.
 
-        Removes only mrope_position_deltas while keeping all other multimodal data
-        (embeddings, images, etc.) needed for context phase processing.
+        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
+
+        Raises:
+            ValueError: If element is not supported or device is invalid
+            RuntimeError: If device transfer fails
         """
-        if not (self.multimodal_data
-                and 'mrope_config' in self.multimodal_data):
-            return
+        supported_elements = ["multimodal_data", "multimodal_input"]
+        self._validate_element(element, supported_elements)
 
-        mrope_config = self.multimodal_data['mrope_config']
-        if 'mrope_position_deltas' in mrope_config:
-            del mrope_config['mrope_position_deltas']
+        data = getattr(self, element)
+        if data is None:
+            return  # Nothing to move
 
-            # Clean up empty mrope_config
-            if not mrope_config:
-                del self.multimodal_data['mrope_config']
+        transformed_data = self._apply_tensor_operation(data,
+                                                        "to_device",
+                                                        device=device,
+                                                        pin_memory=pin_memory)
+        setattr(self, element, transformed_data)
 
     def strip_for_generation(self) -> None:
         """Strip multimodal data for generation processing.
 
-        Keeps only mrope_position_deltas and removes all other multimodal data
+        Keeps only mrope_config 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
+        # Extract mrope_config before clearing
+        mrope_config = None
         if 'mrope_config' in self.multimodal_data:
             mrope_config = self.multimodal_data['mrope_config']
-            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
+        # Clear all data and restore only mrope_config 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 mrope_config is not None:
+            self.multimodal_data['mrope_config'] = mrope_config
 
     def has_content(self) -> bool:
         """Check if this object contains any multimodal data."""
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index 12bb079eaf5a..4a750fb5c244 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -410,6 +410,8 @@ def generate_async(
                 # Only pass it if it has content
                 if not multimodal_params.has_content():
                     multimodal_params = None
+                else:
+                    multimodal_params.to_shared_tensor("multimodal_data")
         else:
             raise TypeError(
                 f"The inputs must be type str or list of int, but got {type(inputs)}"

From 87b8531b8fa0ad8e6393fe54313f1fc0bf29a5aa Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 24 Jul 2025 14:48:32 +0900
Subject: [PATCH 2/9] address review comments

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         |  1 -
 tensorrt_llm/executor/worker.py               | 13 ++----------
 tensorrt_llm/inputs/multimodal.py             | 21 ++++++++++++-------
 3 files changed, 16 insertions(+), 19 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index 7fcf1924f1df..b8788d125e15 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1229,7 +1229,6 @@ def _prepare_tp_inputs(
                 multimodal_runtime=py_multimodal_runtime)
 
             if multimodal_params.has_content():
-                multimodal_params.from_shared_tensor("multimodal_data")
                 multimodal_params.to_device("multimodal_data",
                                             "cuda",
                                             pin_memory=True)
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index db8d84fcc898..b049b2b530a0 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -502,17 +502,8 @@ def _deduce_max_tokens(request: GenerationRequest,
 
             if self._is_pytorch_backend and request.multimodal_params is not None:
                 if request.multimodal_params.multimodal_data is not None:
-                    # Convert back to tensor, as opposite to `to_handle` in `llm.generate_async`
-                    # for values with non-selected keys, it's no-op
-                    request.multimodal_params.to_tensor(
-                        "multimodal_data", key="multimodal_embedding")
-                    embedding = request.multimodal_params.multimodal_data.get(
-                        "multimodal_embedding")
-                    if embedding is not None and embedding.is_cuda:
-                        # make sure the embedding resides on the local device
-                        request.multimodal_params.multimodal_data[
-                            "multimodal_embedding"] = embedding.to("cuda")
-
+                    request.multimodal_params.from_shared_tensor(
+                        "multimodal_data")
                     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 fd647c6fa53e..f6761bc7896c 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -159,6 +159,10 @@ class MultimodalParams:
         multimodal_input: Multimodal input data with hashing information.
         multimodal_data: Processed multimodal data containing embeddings, configurations,
                         and modality-specific data organized by type.
+        multimodal_runtime: Runtime data for tracking multimodal token caching and reuse
+                           during KV cache scenarios. Contains information about cached
+                           tokens, multimodal token positions, and lengths for efficient
+                           processing during inference.
 
     Structure of multimodal_data:
         {
@@ -209,10 +213,15 @@ def _is_shared_tensor_dict(self, obj: Any) -> bool:
 
         # Additional validation based on method_key
         method_key = obj.get('method_key')
-        if method_key == 1:  # CUDA tensor
+
+        # Import here to avoid circular imports
+        from tensorrt_llm._torch.shared_tensor import \
+            _SharedTensorRebuildMethodRegistry
+
+        if method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CUDA:
             cuda_keys = {'tensor_size', 'storage_handle', 'storage_device'}
             return cuda_keys.issubset(obj.keys())
-        elif method_key == 2:  # CPU tensor
+        elif method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU:
             cpu_keys = {'tensor_size', 'storage_handle', 'manager_handle'}
             return cpu_keys.issubset(obj.keys())
 
@@ -253,6 +262,7 @@ def _apply_tensor_operation(
                     input_data):
                 # Convert shared tensor dict back to tensor
                 try:
+                    # Import here to avoid circular imports
                     from tensorrt_llm._torch.shared_tensor import \
                         SharedTensorContainer
 
@@ -275,6 +285,7 @@ def _apply_tensor_operation(
         if isinstance(input_data, torch.Tensor):
             if operation == "to_shared_tensor":
                 try:
+                    # Import here to avoid circular imports
                     from tensorrt_llm._torch.shared_tensor import \
                         SharedTensorContainer
                     return SharedTensorContainer.from_tensor(
@@ -340,14 +351,11 @@ def to_shared_tensor(self, element: str) -> None:
                                                         "to_shared_tensor")
         setattr(self, element, transformed_data)
 
-    def from_shared_tensor(self, element: str, mode: str = "context") -> None:
+    def from_shared_tensor(self, element: str) -> None:
         """Move specified multimodal data element from shared tensor.
 
         Args:
             element: Element to restore ("multimodal_data" or "multimodal_input")
-            mode: Recovery mode - "context" or "generation"
-                - "context": Recover all shared tensors except mrope_position_deltas
-                - "generation": Recover only mrope_position_deltas
 
         Raises:
             ValueError: If element or mode is not supported
@@ -363,7 +371,6 @@ def from_shared_tensor(self, element: str, mode: str = "context") -> None:
 
         restored_data = self._apply_tensor_operation(data, "from_shared_tensor")
         setattr(self, element, restored_data)
-        # print(f"Restored {element} from shared tensor, restored_data: {restored_data}")
 
     def to_device(self,
                   element: str,

From 22c6384464a0370731a862729902c639db1349a1 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 24 Jul 2025 14:50:38 +0900
Subject: [PATCH 3/9] remove print

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

diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index f6761bc7896c..8b2bed1eb2a6 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -266,7 +266,6 @@ def _apply_tensor_operation(
                     from tensorrt_llm._torch.shared_tensor import \
                         SharedTensorContainer
 
-                    # print(f"input_data: {input_data}")
                     return SharedTensorContainer.from_dict(
                         input_data).get_local_view()
                 except Exception as e:

From f1fd7a0f52f609db0a02c5e626cb90855c4e1058 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 24 Jul 2025 15:33:40 +0900
Subject: [PATCH 4/9] address docstring

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

diff --git a/tensorrt_llm/inputs/multimodal.py b/tensorrt_llm/inputs/multimodal.py
index 8b2bed1eb2a6..4cc7a72f9c47 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -357,7 +357,7 @@ def from_shared_tensor(self, element: str) -> None:
             element: Element to restore ("multimodal_data" or "multimodal_input")
 
         Raises:
-            ValueError: If element or mode is not supported
+            ValueError: If element is not supported
             RuntimeError: If tensor restoration fails
         """
         supported_elements = ["multimodal_data", "multimodal_input"]

From a8ff7754023dcb88257e1ed70ee8df8f4f5a47ac Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Thu, 31 Jul 2025 13:21:14 +0900
Subject: [PATCH 5/9] rebase compatibility

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         |  6 ---
 tensorrt_llm/executor/worker.py               |  7 ++-
 tensorrt_llm/inputs/multimodal.py             | 29 ++++++------
 tensorrt_llm/llmapi/llm.py                    |  7 +--
 .../multimodal/test_share_multiparams.py      | 44 ++++++++++++++++---
 5 files changed, 59 insertions(+), 34 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index b8788d125e15..f6fdef3be539 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1229,12 +1229,6 @@ def _prepare_tp_inputs(
                 multimodal_runtime=py_multimodal_runtime)
 
             if multimodal_params.has_content():
-                multimodal_params.to_device("multimodal_data",
-                                            "cuda",
-                                            pin_memory=True)
-                # re-assign the recovered SharedTensorContainer to the request for generation usage
-                setattr(request, "py_multimodal_data",
-                        multimodal_params.multimodal_data)
                 multimodal_params_list.append(multimodal_params)
 
             request.py_batch_idx = request.py_seq_slot
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index b049b2b530a0..940f65919ed2 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -502,8 +502,11 @@ def _deduce_max_tokens(request: GenerationRequest,
 
             if self._is_pytorch_backend and request.multimodal_params is not None:
                 if request.multimodal_params.multimodal_data is not None:
-                    request.multimodal_params.from_shared_tensor(
-                        "multimodal_data")
+                    request.multimodal_params.to_tensor("multimodal_data")
+                    # make sure the tensors reside on the local device
+                    request.multimodal_params.to_device("multimodal_data",
+                                                        "cuda",
+                                                        pin_memory=True)
                     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 4cc7a72f9c47..bc07e0cf78de 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -233,8 +233,8 @@ def _apply_tensor_operation(
         """Apply tensor operations recursively to nested data structures.
 
         This method handles three types of operations:
-        - "to_shared_tensor": Convert tensors to shared tensor dictionaries
-        - "from_shared_tensor": Convert shared tensor dictionaries back to tensors
+        - "to_handle": Convert tensors to shared tensor dictionaries
+        - "to_tensor": Convert shared tensor dictionaries back to tensors
         - "to_device": Move tensors to specified device
 
         Args:
@@ -258,7 +258,7 @@ def _apply_tensor_operation(
 
         # Handle dictionary case
         if isinstance(input_data, dict):
-            if operation == "from_shared_tensor" and self._is_shared_tensor_dict(
+            if operation == "to_tensor" and self._is_shared_tensor_dict(
                     input_data):
                 # Convert shared tensor dict back to tensor
                 try:
@@ -282,7 +282,7 @@ def _apply_tensor_operation(
 
         # Handle tensor case
         if isinstance(input_data, torch.Tensor):
-            if operation == "to_shared_tensor":
+            if operation == "to_handle":
                 try:
                     # Import here to avoid circular imports
                     from tensorrt_llm._torch.shared_tensor import \
@@ -329,38 +329,37 @@ def _validate_element(self, element: str,
                 f"Supported elements: {', '.join(repr(e) for e in supported_elements)}"
             )
 
-    def to_shared_tensor(self, element: str) -> None:
+    def to_handle(self, element: str) -> None:
         """Move specified multimodal data element to shared tensor.
 
         Args:
-            element: Element to move ("multimodal_data" or "multimodal_input")
+            element: Element to move ("multimodal_data")
 
         Raises:
             ValueError: If element is not supported
             RuntimeError: If tensor conversion fails
         """
-        supported_elements = ["multimodal_data", "multimodal_input"]
+        supported_elements = ["multimodal_data"]
         self._validate_element(element, supported_elements)
 
         data = getattr(self, element)
         if data is None:
             return  # Nothing to convert
 
-        transformed_data = self._apply_tensor_operation(data,
-                                                        "to_shared_tensor")
+        transformed_data = self._apply_tensor_operation(data, "to_handle")
         setattr(self, element, transformed_data)
 
-    def from_shared_tensor(self, element: str) -> None:
+    def to_tensor(self, element: str) -> None:
         """Move specified multimodal data element from shared tensor.
 
         Args:
-            element: Element to restore ("multimodal_data" or "multimodal_input")
+            element: Element to restore ("multimodal_data")
 
         Raises:
             ValueError: If element is not supported
             RuntimeError: If tensor restoration fails
         """
-        supported_elements = ["multimodal_data", "multimodal_input"]
+        supported_elements = ["multimodal_data"]
 
         self._validate_element(element, supported_elements)
 
@@ -368,7 +367,7 @@ def from_shared_tensor(self, element: str) -> None:
         if data is None:
             return  # Nothing to restore
 
-        restored_data = self._apply_tensor_operation(data, "from_shared_tensor")
+        restored_data = self._apply_tensor_operation(data, "to_tensor")
         setattr(self, element, restored_data)
 
     def to_device(self,
@@ -378,7 +377,7 @@ def to_device(self,
         """Move specified multimodal data element to target device.
 
         Args:
-            element: Element to move ("multimodal_data" or "multimodal_input")
+            element: Element to move ("multimodal_data")
             device: Target device (e.g., "cuda", "cpu")
             pin_memory: Whether to pin memory for faster transfers
 
@@ -386,7 +385,7 @@ def to_device(self,
             ValueError: If element is not supported or device is invalid
             RuntimeError: If device transfer fails
         """
-        supported_elements = ["multimodal_data", "multimodal_input"]
+        supported_elements = ["multimodal_data"]
         self._validate_element(element, supported_elements)
 
         data = getattr(self, element)
diff --git a/tensorrt_llm/llmapi/llm.py b/tensorrt_llm/llmapi/llm.py
index 4a750fb5c244..78eb0304b5fc 100644
--- a/tensorrt_llm/llmapi/llm.py
+++ b/tensorrt_llm/llmapi/llm.py
@@ -403,15 +403,12 @@ def generate_async(
                         'multimodal_input'),
                     multimodal_data=extra_processed_inputs.get(
                         'multimodal_data'))
-                # Convert to shared tensor handle to reduce IPC overhead
-                # for values with non-selected keys, it's no-op
-                multimodal_params.to_handle("multimodal_data",
-                                            key="multimodal_embedding")
                 # Only pass it if it has content
                 if not multimodal_params.has_content():
                     multimodal_params = None
                 else:
-                    multimodal_params.to_shared_tensor("multimodal_data")
+                    # Convert to shared tensor handle to reduce IPC overhead
+                    multimodal_params.to_handle("multimodal_data")
         else:
             raise TypeError(
                 f"The inputs must be type str or list of int, but got {type(inputs)}"
diff --git a/tests/unittest/_torch/multimodal/test_share_multiparams.py b/tests/unittest/_torch/multimodal/test_share_multiparams.py
index d4ce40f6332c..228adf948d25 100644
--- a/tests/unittest/_torch/multimodal/test_share_multiparams.py
+++ b/tests/unittest/_torch/multimodal/test_share_multiparams.py
@@ -39,14 +39,46 @@ def test_to_handle_none_multimodal_data(self):
         params.to_handle("multimodal_data")
         self.assertEqual(params.multimodal_data, {})
 
+    def test_to_handle_unsupported_element(self):
+        """Test to_handle raises ValueError for unsupported elements."""
         params = MultimodalParams()
         multimodal_input = MultimodalInput(
             multimodal_hashes=[[1, 2, 3, 4, 5, 6, 7, 8]] * 2,
             multimodal_positions=[0, 10],
             multimodal_lengths=[2, 2])
         params.multimodal_input = multimodal_input
-        params.to_handle("multimodal_input")
-        self.assertEqual(params.multimodal_input, multimodal_input)
+
+        with self.assertRaises(ValueError) as context:
+            params.to_handle("multimodal_input")
+
+        self.assertIn("Unsupported element 'multimodal_input'",
+                      str(context.exception))
+        self.assertIn("Supported elements: 'multimodal_data'",
+                      str(context.exception))
+
+    def test_to_tensor_unsupported_element(self):
+        """Test to_tensor raises ValueError for unsupported elements."""
+        params = MultimodalParams()
+
+        with self.assertRaises(ValueError) as context:
+            params.to_tensor("multimodal_input")
+
+        self.assertIn("Unsupported element 'multimodal_input'",
+                      str(context.exception))
+        self.assertIn("Supported elements: 'multimodal_data'",
+                      str(context.exception))
+
+    def test_to_device_unsupported_element(self):
+        """Test to_device raises ValueError for unsupported elements."""
+        params = MultimodalParams()
+
+        with self.assertRaises(ValueError) as context:
+            params.to_device("multimodal_input", device="cuda", pin_memory=True)
+
+        self.assertIn("Unsupported element 'multimodal_input'",
+                      str(context.exception))
+        self.assertIn("Supported elements: 'multimodal_data'",
+                      str(context.exception))
 
     def test_to_tensor_basic_handle(self):
         """Test converting a basic handle back to tensor."""
@@ -54,9 +86,9 @@ def test_to_tensor_basic_handle(self):
         params.multimodal_data = {"multimodal_embedding": self.mm_embedding}
 
         # Convert to handle
-        params.to_handle("multimodal_data", key="multimodal_embedding")
+        params.to_handle("multimodal_data")
         # Convert back to tensor
-        params.to_tensor("multimodal_data", key="multimodal_embedding")
+        params.to_tensor("multimodal_data")
 
         result = params.multimodal_data["multimodal_embedding"]
         self.assertIsInstance(result, torch.Tensor)
@@ -67,8 +99,8 @@ def test_to_tensor_all_handles(self):
         params = MultimodalParams()
         params.multimodal_data = self.sample_multimodal_data.copy()
 
-        params.to_handle("multimodal_data", key=None)
-        params.to_tensor("multimodal_data", key=None)
+        params.to_handle("multimodal_data")
+        params.to_tensor("multimodal_data")
 
         self.assertTrue(
             torch.allclose(params.multimodal_data["multimodal_embedding"],

From 3ec4f916573469784c10ae676dff90bb44d25d4c Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Mon, 4 Aug 2025 13:31:22 +0900
Subject: [PATCH 6/9] address comments

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../_torch/pyexecutor/model_engine.py         | 10 ++--
 tensorrt_llm/executor/worker.py               |  4 --
 tensorrt_llm/inputs/multimodal.py             | 48 +++++++------------
 3 files changed, 25 insertions(+), 37 deletions(-)

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index f6fdef3be539..c3bf6d8b3c2b 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1229,6 +1229,11 @@ def _prepare_tp_inputs(
                 multimodal_runtime=py_multimodal_runtime)
 
             if multimodal_params.has_content():
+                multimodal_params.to_device("multimodal_data",
+                                            "cuda",
+                                            pin_memory=True)
+                #re-assign the multimodal_data to the request after to_device for generation requests
+                request.py_multimodal_data = multimodal_params.multimodal_data
                 multimodal_params_list.append(multimodal_params)
 
             request.py_batch_idx = request.py_seq_slot
@@ -1263,9 +1268,8 @@ def _prepare_tp_inputs(
                 multimodal_data=request.py_multimodal_data)
             if multimodal_params.has_content():
                 multimodal_params.strip_for_generation()
-                # re-assign the recovered SharedTensorContainer to the request after strip_for_generation
-                setattr(request, "py_multimodal_data",
-                        multimodal_params.multimodal_data)
+                # re-assign the multimodal_data to the request after strip_for_generation for another generation request
+                request.py_multimodal_data = multimodal_params.multimodal_data
                 multimodal_params_list.append(multimodal_params)
         extend_requests += extend_dummy_requests
 
diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index 940f65919ed2..160975e92e60 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -503,10 +503,6 @@ def _deduce_max_tokens(request: GenerationRequest,
             if self._is_pytorch_backend and request.multimodal_params is not None:
                 if request.multimodal_params.multimodal_data is not None:
                     request.multimodal_params.to_tensor("multimodal_data")
-                    # make sure the tensors reside on the local device
-                    request.multimodal_params.to_device("multimodal_data",
-                                                        "cuda",
-                                                        pin_memory=True)
                     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 bc07e0cf78de..77eec0ff6a66 100644
--- a/tensorrt_llm/inputs/multimodal.py
+++ b/tensorrt_llm/inputs/multimodal.py
@@ -312,35 +312,20 @@ def _apply_tensor_operation(
         # For any other type, return as-is
         return input_data
 
-    def _validate_element(self, element: str,
-                          supported_elements: List[str]) -> None:
-        """Validate that the element is supported.
-
-        Args:
-            element: Element name to validate
-            supported_elements: List of supported element names
-
-        Raises:
-            ValueError: If element is not supported
-        """
-        if element not in supported_elements:
-            raise ValueError(
-                f"Unsupported element '{element}'. "
-                f"Supported elements: {', '.join(repr(e) for e in supported_elements)}"
-            )
-
     def to_handle(self, element: str) -> None:
         """Move specified multimodal data element to shared tensor.
 
         Args:
-            element: Element to move ("multimodal_data")
+            element: Element to move (only "multimodal_data" is supported)
 
         Raises:
-            ValueError: If element is not supported
+            ValueError: If element is not "multimodal_data"
             RuntimeError: If tensor conversion fails
         """
-        supported_elements = ["multimodal_data"]
-        self._validate_element(element, supported_elements)
+        if element != "multimodal_data":
+            raise ValueError(
+                f"Unsupported element '{element}'. Only 'multimodal_data' is supported."
+            )
 
         data = getattr(self, element)
         if data is None:
@@ -353,15 +338,16 @@ def to_tensor(self, element: str) -> None:
         """Move specified multimodal data element from shared tensor.
 
         Args:
-            element: Element to restore ("multimodal_data")
+            element: Element to restore (only "multimodal_data" is supported)
 
         Raises:
-            ValueError: If element is not supported
+            ValueError: If element is not "multimodal_data"
             RuntimeError: If tensor restoration fails
         """
-        supported_elements = ["multimodal_data"]
-
-        self._validate_element(element, supported_elements)
+        if element != "multimodal_data":
+            raise ValueError(
+                f"Unsupported element '{element}'. Only 'multimodal_data' is supported."
+            )
 
         data = getattr(self, element)
         if data is None:
@@ -377,16 +363,18 @@ def to_device(self,
         """Move specified multimodal data element to target device.
 
         Args:
-            element: Element to move ("multimodal_data")
+            element: Element to move (only "multimodal_data" is supported)
             device: Target device (e.g., "cuda", "cpu")
             pin_memory: Whether to pin memory for faster transfers
 
         Raises:
-            ValueError: If element is not supported or device is invalid
+            ValueError: If element is not "multimodal_data" or device is invalid
             RuntimeError: If device transfer fails
         """
-        supported_elements = ["multimodal_data"]
-        self._validate_element(element, supported_elements)
+        if element != "multimodal_data":
+            raise ValueError(
+                f"Unsupported element '{element}'. Only 'multimodal_data' is supported."
+            )
 
         data = getattr(self, element)
         if data is None:

From 72d39b0a3faba31538c9cd117da246817b09a3bb Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Mon, 4 Aug 2025 14:11:18 +0900
Subject: [PATCH 7/9] add one-line explanation

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

diff --git a/tensorrt_llm/executor/worker.py b/tensorrt_llm/executor/worker.py
index 160975e92e60..2d8c1c0255d1 100644
--- a/tensorrt_llm/executor/worker.py
+++ b/tensorrt_llm/executor/worker.py
@@ -486,7 +486,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`.
+                # 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=(
@@ -502,6 +502,7 @@ def _deduce_max_tokens(request: GenerationRequest,
 
             if self._is_pytorch_backend and request.multimodal_params is not None:
                 if request.multimodal_params.multimodal_data is not None:
+                    # NOTE: Deserialize SharedTensor handle to actual tensor
                     request.multimodal_params.to_tensor("multimodal_data")
                     executor_request.py_multimodal_data = request.multimodal_params.multimodal_data
 

From 8a91c216d28153ecb28939f00ac3b77ca477bd24 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Tue, 5 Aug 2025 11:09:54 +0900
Subject: [PATCH 8/9] move strip_for_generation

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

diff --git a/tensorrt_llm/_torch/pyexecutor/model_engine.py b/tensorrt_llm/_torch/pyexecutor/model_engine.py
index c3bf6d8b3c2b..f98ba385ca7f 100644
--- a/tensorrt_llm/_torch/pyexecutor/model_engine.py
+++ b/tensorrt_llm/_torch/pyexecutor/model_engine.py
@@ -1266,9 +1266,12 @@ def _prepare_tp_inputs(
             # Multimodal
             multimodal_params = MultimodalParams(
                 multimodal_data=request.py_multimodal_data)
+            multimodal_params.strip_for_generation()
             if multimodal_params.has_content():
-                multimodal_params.strip_for_generation()
-                # re-assign the multimodal_data to the request after strip_for_generation for another generation request
+                multimodal_params.to_device("multimodal_data",
+                                            "cuda",
+                                            pin_memory=True)
+                # re-assign the multimodal_data to the request after strip_for_generation for another generation request,
                 request.py_multimodal_data = multimodal_params.multimodal_data
                 multimodal_params_list.append(multimodal_params)
         extend_requests += extend_dummy_requests

From 204c30aa08c25b12ba76f2f5e7dc722840a96e40 Mon Sep 17 00:00:00 2001
From: yechank <161688079+yechank-nvidia@users.noreply.github.com>
Date: Fri, 8 Aug 2025 10:43:34 +0900
Subject: [PATCH 9/9] modify test

Signed-off-by: yechank <161688079+yechank-nvidia@users.noreply.github.com>
---
 .../multimodal/test_share_multiparams.py      | 77 ++++++++++++-------
 1 file changed, 51 insertions(+), 26 deletions(-)

diff --git a/tests/unittest/_torch/multimodal/test_share_multiparams.py b/tests/unittest/_torch/multimodal/test_share_multiparams.py
index 228adf948d25..343c2f537212 100644
--- a/tests/unittest/_torch/multimodal/test_share_multiparams.py
+++ b/tests/unittest/_torch/multimodal/test_share_multiparams.py
@@ -53,32 +53,6 @@ def test_to_handle_unsupported_element(self):
 
         self.assertIn("Unsupported element 'multimodal_input'",
                       str(context.exception))
-        self.assertIn("Supported elements: 'multimodal_data'",
-                      str(context.exception))
-
-    def test_to_tensor_unsupported_element(self):
-        """Test to_tensor raises ValueError for unsupported elements."""
-        params = MultimodalParams()
-
-        with self.assertRaises(ValueError) as context:
-            params.to_tensor("multimodal_input")
-
-        self.assertIn("Unsupported element 'multimodal_input'",
-                      str(context.exception))
-        self.assertIn("Supported elements: 'multimodal_data'",
-                      str(context.exception))
-
-    def test_to_device_unsupported_element(self):
-        """Test to_device raises ValueError for unsupported elements."""
-        params = MultimodalParams()
-
-        with self.assertRaises(ValueError) as context:
-            params.to_device("multimodal_input", device="cuda", pin_memory=True)
-
-        self.assertIn("Unsupported element 'multimodal_input'",
-                      str(context.exception))
-        self.assertIn("Supported elements: 'multimodal_data'",
-                      str(context.exception))
 
     def test_to_tensor_basic_handle(self):
         """Test converting a basic handle back to tensor."""
@@ -122,5 +96,56 @@ def test_to_tensor_all_handles(self):
                          self.image["image_width"])
 
 
+class TestMultimodalParamsDeviceTransfer(unittest.TestCase):
+    """Test cases for to_device method in MultimodalParams."""
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.mm_embedding = torch.randn(3, 4, 5)
+        self.mrope_config = {
+            "mrope_rotary_cos_sin": torch.randn(2, 3),
+            "mrope_position_deltas": torch.randn(5),
+        }
+        self.image = {
+            "pixel_values": torch.randn(1, 3, 224, 224),
+            "image_height": [224],
+            "image_width": [224],
+        }
+        self.sample_multimodal_data = {
+            "multimodal_embedding": self.mm_embedding,
+            "mrope_config": self.mrope_config,
+            "image": self.image,
+        }
+
+    def test_to_device_basic(self):
+        """Test converting a basic data to device."""
+        params = MultimodalParams()
+        params.multimodal_data = {"multimodal_embedding": self.mm_embedding}
+
+        params.to_device("multimodal_data", device="cuda:0", pin_memory=True)
+
+        result = params.multimodal_data["multimodal_embedding"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+    def test_to_device_all_data(self):
+        """Test converting all data to device."""
+        params = MultimodalParams()
+        params.multimodal_data = self.sample_multimodal_data.copy()
+
+        params.to_device("multimodal_data", device="cuda:0", pin_memory=True)
+
+        result = params.multimodal_data["multimodal_embedding"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+        result = params.multimodal_data["mrope_config"]["mrope_rotary_cos_sin"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+        result = params.multimodal_data["mrope_config"]["mrope_position_deltas"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+        result = params.multimodal_data["image"]["pixel_values"]
+        self.assertEqual(result.device, torch.device("cuda:0"))
+
+
 if __name__ == "__main__":
     unittest.main()