From a7730bf1348c18bf64f378c0693974ce386aac5f Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sat, 21 Jun 2025 00:44:57 -0700
Subject: [PATCH 1/7] Add share tensor utility

---
 tensorrt_llm/_torch/multimodal/mm_utils.py    | 355 ++++++++++++++++++
 .../_torch/multimodal/test_share_tensor.py    | 145 +++++++
 2 files changed, 500 insertions(+)
 create mode 100644 tensorrt_llm/_torch/multimodal/mm_utils.py
 create mode 100644 tests/unittest/_torch/multimodal/test_share_tensor.py

diff --git a/tensorrt_llm/_torch/multimodal/mm_utils.py b/tensorrt_llm/_torch/multimodal/mm_utils.py
new file mode 100644
index 000000000000..f3d87b98bdea
--- /dev/null
+++ b/tensorrt_llm/_torch/multimodal/mm_utils.py
@@ -0,0 +1,355 @@
+import logging
+import base64
+from typing import Any, Callable, Dict, Tuple
+
+import torch
+from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor, reduce_storage, rebuild_tensor, rebuild_meta_tensor, rebuild_storage_filename, rebuild_typed_storage
+from torch.multiprocessing import get_sharing_strategy, set_sharing_strategy
+
+logger = logging.getLogger(__name__)
+
+
+class _SharedTensorRebuildMethodRegistry:
+    """Registry for tensor rebuild methods with fixed keys for common methods.
+    
+    This registry is used by the SharedTensorContainer to manage PyTorch tensor
+    rebuild methods on a consumer process.
+    
+    This class maintains a mapping of numeric keys to rebuild methods.
+    Common methods are pre-registered with fixed keys for consistency.
+    """
+    # Fixed keys for common rebuild methods
+    REBUILD_CUDA = 1
+    REBUILD_CPU = 2
+    REBUILD_META = 3
+
+    _registry: Dict[int, Callable] = {}
+
+    @classmethod
+    def initialize(cls):
+        """Initialize the registry with common rebuild methods."""
+        # Register common methods with fixed keys
+        cls._registry[cls.REBUILD_CUDA] = rebuild_cuda_tensor
+        cls._registry[cls.REBUILD_CPU] = rebuild_tensor
+        cls._registry[cls.REBUILD_META] = rebuild_meta_tensor
+
+    @classmethod
+    def register(cls, method: Callable) -> int:
+        """Register a rebuild method and return its key.
+
+        Args:
+            method: The rebuild method to register
+
+        Returns:
+            The numeric key assigned to the method
+        """
+        if method == rebuild_cuda_tensor:
+            return cls.REBUILD_CUDA
+        if method == rebuild_tensor:
+            return cls.REBUILD_CPU
+        if method == rebuild_meta_tensor:
+            # TODO: add support for meta tensor if needed
+            return cls.REBUILD_META
+        raise NotImplementedError("Other rebuild methods are not supported yet")
+
+    @classmethod
+    def get_method(cls, key: int) -> Callable:
+        """Get a rebuild method by its key.
+
+        Args:
+            key: The numeric key of the method
+
+        Returns:
+            The registered rebuild method
+
+        Raises:
+            KeyError: If the key is not found in the registry
+        """
+        if key not in cls._registry:
+            raise KeyError(f"No rebuild method registered with key {key}")
+        return cls._registry[key]
+
+
+class SharedTensorContainer:
+    """A class for sharing tensors between processes (non-Python multiprocessing processes).
+    
+    This is an intermediate solution to accommodate communication between two independent processes
+    that are not spawned by Python multiprocessing. It uses shared memory or CUDA IPC to avoid
+    serialization/IPC costs when communicating tensors between processes.
+    
+    Key Features:
+    - Uses PyTorch's reduction methods (reduce_tensor, rebuild_tensor, etc.) to handle tensor sharing
+    - Supports both CPU tensors (via shared memory) and CUDA tensors (via CUDA IPC)
+    - Provides serialization capabilities for cross-process communication
+    - Avoids expensive tensor serialization/deserialization/IPC overhead
+    
+    Architecture:
+    - Producer process: Creates tensor -> reduces -> convert handle to dict -> sends via IPC
+    - Consumer process: Receives dict -> convert dict to handle -> rebuilds tensor -> gets local view
+    
+    Note: This is a temporary solution. In the future, if we embrace more Python environment
+    or PyTorch orchestration methods (like torch.multiprocessing, Ray, etc.), this intermediate
+    layer would not be needed as those frameworks provide native tensor sharing capabilities.
+    
+    Note: Whenever you call reduce_tensor, you must call the corresponding rebuild method at 
+    consumer process(es), otherwise, the producer process cannot release the memory to caching 
+    allocator as the inner refcount never reaches zero.
+    """
+    def __init__(self, method_key: int, tensor_handle: Dict[str, Any]):
+        """Initialize the SharedTensorContainer.
+        
+        Args:
+            method_key: Registry key for the rebuild method (CUDA, CPU, etc.)
+            tensor_handle: Tensor handle that can be used to rebuild the tensor on a consumer process
+        """
+        self.method_key = method_key
+        self.tensor_handle = tensor_handle
+
+    @staticmethod
+    def cuda_handle_to_dict(tensor_handle) -> Dict[str, Any]:
+        """Convert CUDA tensor handle to serializable dictionary for IPC
+        
+        This method converts PyTorch's CUDA tensor reduction (cudaIPC) handle into a format that can be
+        safely serialized and transmitted between any two processes. It handles binary data by 
+        encoding it in base64 to ensure JSON compatibility.
+        
+        The CUDA handle contains references to GPU memory, CUDA events, and reference counters
+        that need to be properly serialized for cross-process sharing via CUDA IPC.
+
+        Returns:
+            Dictionary containing the serialized CUDA tensor information
+
+        Raises:
+            KeyError: If required tensor information is missing
+            ValueError: If tensor information cannot be serialized
+        """
+        try:
+            # tensor_handle is a tuple returned by reduce_tensor
+            tensor_info = tensor_handle
+            # Convert tensor info to a basic dict with only serializable values
+            serializable_info = {
+                # tensor_info[0] is the type of the tensor, which is "torch.Tensor"
+                "tensor_size": list(tensor_info[1]),
+                "tensor_stride": list(tensor_info[2]),
+                "tensor_offset": tensor_info[3],
+                "dtype": str(tensor_info[5]),
+                "storage_device": tensor_info[6],
+                "storage_handle": base64.b64encode(tensor_info[7]).decode('utf-8'),
+                "storage_size_bytes": tensor_info[8],
+                "storage_offset_bytes": tensor_info[9],
+                "requires_grad": tensor_info[10],
+                "ref_counter_handle": base64.b64encode(tensor_info[11]).decode('utf-8'),
+                "ref_counter_offset": tensor_info[12],
+                "event_handle": base64.b64encode(tensor_info[13]).decode('utf-8'),
+                "event_sync_required": tensor_info[14]
+            }
+            return serializable_info
+        except IndexError as e:
+            raise KeyError(f"Missing required tensor information: {e}")
+        except Exception as e:
+            raise ValueError(f"Failed to serialize tensor information: {e}")
+
+    @staticmethod
+    def cpu_handle_to_dict(meta_data: Dict[str, Any], storage_metadata: Dict[str, Any]) -> Dict[str, Any]:
+        """Convert CPU tensor handle to serializable dictionary for IPC
+        
+        This method converts PyTorch's CPU tensor reduction handle into a format that can be
+        safely serialized and transmitted between any two processes. For CPU tensors, we use
+        shared memory via file_system sharing strategy as fd strategy is not serializable.
+        
+        Args:
+            meta_data: Tensor metadata (size, stride, offset, etc.)
+            storage_metadata: Storage metadata (handle, size, dtype, etc.)
+            
+        Returns:
+            Dictionary containing the serialized CPU tensor information
+        """
+        try:
+            serializable_info = {
+                "tensor_storage_offset": meta_data[0],
+                "tensor_size": list(meta_data[1]),
+                "tensor_stride": list(meta_data[2]),
+                "manager_handle": base64.b64encode(storage_metadata[0]).decode('utf-8'),
+                "storage_handle": base64.b64encode(storage_metadata[1]).decode('utf-8'),
+                "storage_size": storage_metadata[2],
+                "storage_dtype": str(storage_metadata[3])
+            }
+            return serializable_info
+        except IndexError as e:
+            raise KeyError(f"Missing required tensor information: {e}")
+        except Exception as e:
+            raise ValueError(f"Failed to serialize tensor information: {e}")
+
+
+    @staticmethod
+    def dict_to_cuda_handle(tensor_info: Dict[str, Any]) -> Tuple:
+        """Reconstruct CUDA tensor handle from serialized dictionary
+        
+        This method reconstructs a CUDA tensor handle from a previously serialized dictionary
+        that was received from another process. 
+        
+        Args:
+            tensor_info: Dictionary containing the serialized CUDA tensor information
+                with the same keys as returned by cuda_handle_to_dict()
+
+        Returns:
+            A tuple representing the CUDA tensor handle for PyTorch's rebuild_cuda_tensor
+        """
+        try:
+            # Decode base64 encoded binary data
+            storage_handle = base64.b64decode(tensor_info['storage_handle'])
+            ref_counter_handle = base64.b64decode(tensor_info['ref_counter_handle'])
+            event_handle = base64.b64decode(tensor_info['event_handle'])
+
+            # Reconstruct the tensor handle
+            tensor_handle = (torch.Tensor,
+                             tuple(tensor_info['tensor_size']),
+                             tuple(tensor_info['tensor_stride']),
+                             tensor_info['tensor_offset'],
+                             torch.storage.TypedStorage,
+                             eval(tensor_info['dtype']),
+                             tensor_info['storage_device'],
+                             storage_handle,
+                             tensor_info['storage_size_bytes'],
+                             tensor_info['storage_offset_bytes'],
+                             tensor_info['requires_grad'],
+                             ref_counter_handle,
+                             tensor_info['ref_counter_offset'],
+                             event_handle,
+                             tensor_info['event_sync_required'])
+
+            return tensor_handle
+        except KeyError as e:
+            raise KeyError(f"Missing required tensor information: {e}")
+        except Exception as e:
+            raise ValueError(f"Failed to deserialize tensor information: {e}")
+
+    @staticmethod
+    def dict_to_cpu_handle(tensor_info: Dict[str, Any]) -> Tuple:
+        """Reconstruct CPU tensor handle from serialized dictionary
+        
+        This method reconstructs a CPU tensor handle from a previously serialized dictionary
+        that was received from another process.
+        
+        The reconstructed handle allows the consumer process to access the same memory
+        region that was shared by the producer process, avoiding data copying.
+
+        Args:
+            tensor_info: Dictionary containing the serialized CPU tensor information
+                with the same keys as returned by cpu_handle_to_dict()
+                
+        Returns:
+            A tuple representing the CPU tensor handle for PyTorch's rebuild_tensor
+        """
+        try:
+            manager_handle = base64.b64decode(tensor_info['manager_handle'])
+            storage_handle = base64.b64decode(tensor_info['storage_handle'])
+            storage_metadata = (torch.storage.TypedStorage, 
+                                manager_handle, 
+                                storage_handle, 
+                                tensor_info['storage_size'],
+                                eval(tensor_info['storage_dtype']))
+            storage = rebuild_storage_filename(*storage_metadata)
+            if not isinstance(storage, torch.storage.TypedStorage):
+                storage = rebuild_typed_storage(storage, eval(tensor_info['storage_dtype']))
+
+            meta_data = (tensor_info['tensor_storage_offset'],
+                         tuple(tensor_info['tensor_size']),
+                         tuple(tensor_info['tensor_stride']),
+                         False) # requires_grad is always False for cpu tensor
+            tensor_handle = (torch.Tensor,
+                             storage,
+                             meta_data)
+            return tensor_handle
+        except KeyError as e:
+            raise KeyError(f"Missing required tensor information: {e}")
+        except Exception as e:
+            raise ValueError(f"Failed to deserialize tensor information: {e}")
+
+    @classmethod
+    def from_tensor(cls, tensor: torch.Tensor) -> 'SharedTensorContainer':
+        """Create a SharedTensorContainer from a local tensor (Producer side).
+        
+        This method is called by the producer process to prepare a tensor for sharing
+        with other processes. It uses PyTorch's reduction methods to create a handle
+        that can be efficiently transmitted and reconstructed by consumer processes.
+        
+        - CUDA tensors: Uses CUDA IPC for GPU memory sharing
+        - CPU tensors: Uses file_system sharing strategy for shared memory
+
+        Args:
+            tensor: The tensor to share
+
+        Returns:
+            SharedTensorContainer instance that can be serialized later for IPC
+        """
+        rebuild_method, tensor_handle = reduce_tensor(tensor)
+        method_key = _SharedTensorRebuildMethodRegistry.register(rebuild_method)
+        return cls(method_key, tensor_handle)
+
+    @classmethod
+    def from_dict(cls, tensor_info: Dict[str, Any]) -> 'SharedTensorContainer':
+        """Create a SharedTensorContainer from a serialized dictionary (Consumer side).
+        
+        This method is called by the consumer process to reconstruct a SharedTensorContainer
+        from serialized data received via IPC.
+
+        Args:
+            tensor_info: Dictionary containing the serialized tensor information
+                received from the producer process via IPC
+                
+        Returns:
+            SharedTensorContainer instance ready for tensor reconstruction
+            
+        Raises:
+            ValueError: If the method_key is not supported
+        """
+        method_key = tensor_info['method_key']
+        if method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CUDA:
+            tensor_handle = SharedTensorContainer.dict_to_cuda_handle(tensor_info)
+        elif method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU:
+            tensor_handle = SharedTensorContainer.dict_to_cpu_handle(tensor_info)
+        else:
+            raise ValueError(f"Unsupported shared tensor method key: {method_key}")
+        return cls(method_key, tensor_handle)
+
+    def get_local_view(self) -> torch.Tensor:
+        """Convert the shared tensor back to a local tensor (Consumer side).
+        
+        This method is called by the consumer process to obtain the actual tensor
+        from the SharedTensorContainer.
+
+        Returns:
+            The reconstructed tensor
+        """
+        rebuild_method = _SharedTensorRebuildMethodRegistry.get_method(self.method_key)
+        return rebuild_method(*self.tensor_handle)
+
+    def dump_to_dict(self) -> Dict[str, Any]:
+        """Convert this container to a dictionary for direct IPC (Producer side).
+        
+        This method is called by the producer process to serialize the SharedTensorContainer instance
+        into a format that is JSON compatible and can be transmitted via IPC. 
+
+        Returns:
+            Dictionary containing the serialized tensor information
+        """
+        if self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CUDA:
+            tensor_dict = SharedTensorContainer.cuda_handle_to_dict(self.tensor_handle)
+        elif self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU:
+            sharing_strategy = get_sharing_strategy()
+            # Here we use file_system sharing strategy to make it serializable between two non-python independent processes
+            set_sharing_strategy("file_system")
+            storage = self.tensor_handle[1] 
+            meta_data = self.tensor_handle[2]
+            storage_handle = reduce_storage(storage)
+            # restore the original sharing strategy
+            set_sharing_strategy(sharing_strategy)
+            # exclude the first element which is the type of the storage
+            storage_metadata = storage_handle[-1][1:] 
+            tensor_dict = SharedTensorContainer.cpu_handle_to_dict(meta_data, storage_metadata)
+        else:
+            raise ValueError(f"Unsupported tensor device: {self.method_key}")
+        
+        tensor_dict["method_key"] = self.method_key
+        return tensor_dict
diff --git a/tests/unittest/_torch/multimodal/test_share_tensor.py b/tests/unittest/_torch/multimodal/test_share_tensor.py
new file mode 100644
index 000000000000..30fe7a94f827
--- /dev/null
+++ b/tests/unittest/_torch/multimodal/test_share_tensor.py
@@ -0,0 +1,145 @@
+import unittest
+import multiprocessing as mp
+import torch
+
+from tensorrt_llm._torch.multimodal.mm_utils import SharedTensorContainer, _SharedTensorRebuildMethodRegistry
+
+
+class TestShareTensor(unittest.TestCase):
+    """Test cases for sharing tensors between processes."""
+
+    @classmethod
+    def setUpClass(cls):
+        """Initialize the registry before running tests."""
+        _SharedTensorRebuildMethodRegistry.initialize()
+
+    def setUp(self):
+        """Set up test fixtures."""
+        self.ref_tensor = torch.randn(3, 4, 5)
+        self.cuda_available = torch.cuda.is_available()
+        if self.cuda_available:
+            torch.cuda.set_device(0)
+
+    @staticmethod
+    def _producer(q, tensor, device=None):
+        """Producer: create tensor and share it."""
+        try:
+            if device is not None:
+                if device == "cuda":
+                    tensor = tensor.cuda()
+                elif device == "cpu":
+                    tensor = tensor.cpu()
+            container = SharedTensorContainer.from_tensor(tensor)
+            q.put(('success', container.dump_to_dict()))
+        except Exception as e:
+            q.put(('error', str(e)))
+
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
+    def test_share_cuda_tensor(self):
+        """Test CUDA tensor sharing between processes."""
+        mp.set_start_method('spawn', force=True)
+        queue = mp.Queue()
+
+        # Producer process
+        producer = mp.Process(target=self._producer, args=(queue, self.ref_tensor, "cuda"))
+        producer.start()
+        status, data = queue.get(timeout=100)
+        # Verify
+        self.assertEqual(status, 'success')
+        reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+        self.assertTrue(torch.allclose(reconstructed.cpu(), self.ref_tensor))
+        del reconstructed
+        producer.join()
+
+    def test_share_cpu_tensor(self):
+        """Test CPU tensor sharing between processes."""
+        mp.set_start_method('spawn', force=True)
+        queue = mp.Queue()
+
+        # Producer process
+        producer = mp.Process(target=self._producer, args=(queue, self.ref_tensor, "cpu"))
+        producer.start()
+        status, data = queue.get(timeout=100)
+        # Verify
+        self.assertEqual(status, 'success')
+        reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+        self.assertTrue(torch.allclose(reconstructed, self.ref_tensor))
+        producer.join()
+
+    def test_share_tensor_different_shapes(self):
+        """Test CPU tensor sharing with different tensor shapes."""
+        mp.set_start_method('spawn', force=True)
+        queue = mp.Queue()
+        test_shapes = [
+            (1,),           
+            (2, 3),         
+            (1, 2, 3, 4),   
+            (10,),          
+        ]
+
+        for shape in test_shapes:
+            with self.subTest(shape=shape):
+                test_tensor = torch.randn(shape)
+                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cpu"))
+                producer.start()
+                status, data = queue.get(timeout=100)
+                
+                self.assertEqual(status, 'success')
+                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+                self.assertTrue(torch.allclose(reconstructed, test_tensor))
+                producer.join()
+            
+            with self.subTest(shape=shape):
+                test_tensor = torch.randn(shape)
+                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cuda"))
+                producer.start()
+                status, data = queue.get(timeout=100)
+                self.assertEqual(status, 'success')
+                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+                self.assertTrue(torch.allclose(reconstructed, test_tensor.cuda()))
+                del reconstructed
+                producer.join()
+
+    def test_share_tensor_different_dtypes(self):
+        """Test CPU tensor sharing with different data types."""
+        mp.set_start_method('spawn', force=True)
+        queue = mp.Queue()
+
+        # Test different data types
+        test_dtypes = [
+            torch.float32,
+            torch.float64,
+            torch.int32,
+            torch.int64,
+        ]
+
+        for dtype in test_dtypes:
+            with self.subTest(dtype=dtype):
+                test_tensor = torch.randn(2, 3).to(dtype)
+                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cpu"))
+                producer.start()
+                status, data = queue.get(timeout=100)
+                
+                self.assertEqual(status, 'success')
+                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+                self.assertTrue(torch.allclose(reconstructed, test_tensor))
+                self.assertEqual(reconstructed.dtype, test_tensor.dtype)
+                producer.join()
+            
+            with self.subTest(dtype=dtype):
+                test_tensor = torch.randn(2, 3).to(dtype)
+                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cuda"))
+                producer.start()
+                status, data = queue.get(timeout=100)
+                self.assertEqual(status, 'success')
+                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+                self.assertTrue(torch.allclose(reconstructed, test_tensor.cuda()))
+                self.assertEqual(reconstructed.dtype, test_tensor.dtype)
+                del reconstructed
+                producer.join()
+
+        
+
+if __name__ == '__main__':
+    mp.set_start_method('spawn', force=True)
+    unittest.main() 
\ No newline at end of file

From f73df16da101efbfeda24a52e1ca422cc3366cde Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sat, 21 Jun 2025 07:18:03 -0700
Subject: [PATCH 2/7] Formatting

---
 tensorrt_llm/_torch/multimodal/mm_utils.py    | 187 ++++++++++--------
 .../_torch/multimodal/test_share_tensor.py    |  61 +++---
 2 files changed, 142 insertions(+), 106 deletions(-)

diff --git a/tensorrt_llm/_torch/multimodal/mm_utils.py b/tensorrt_llm/_torch/multimodal/mm_utils.py
index f3d87b98bdea..de2b7defad36 100644
--- a/tensorrt_llm/_torch/multimodal/mm_utils.py
+++ b/tensorrt_llm/_torch/multimodal/mm_utils.py
@@ -1,20 +1,22 @@
-import logging
 import base64
+import logging
 from typing import Any, Callable, Dict, Tuple
 
 import torch
-from torch.multiprocessing.reductions import rebuild_cuda_tensor, reduce_tensor, reduce_storage, rebuild_tensor, rebuild_meta_tensor, rebuild_storage_filename, rebuild_typed_storage
 from torch.multiprocessing import get_sharing_strategy, set_sharing_strategy
+from torch.multiprocessing.reductions import (
+    rebuild_cuda_tensor, rebuild_meta_tensor, rebuild_storage_filename,
+    rebuild_tensor, rebuild_typed_storage, reduce_storage, reduce_tensor)
 
 logger = logging.getLogger(__name__)
 
 
 class _SharedTensorRebuildMethodRegistry:
     """Registry for tensor rebuild methods with fixed keys for common methods.
-    
+
     This registry is used by the SharedTensorContainer to manage PyTorch tensor
     rebuild methods on a consumer process.
-    
+
     This class maintains a mapping of numeric keys to rebuild methods.
     Common methods are pre-registered with fixed keys for consistency.
     """
@@ -72,32 +74,33 @@ def get_method(cls, key: int) -> Callable:
 
 class SharedTensorContainer:
     """A class for sharing tensors between processes (non-Python multiprocessing processes).
-    
+
     This is an intermediate solution to accommodate communication between two independent processes
     that are not spawned by Python multiprocessing. It uses shared memory or CUDA IPC to avoid
     serialization/IPC costs when communicating tensors between processes.
-    
+
     Key Features:
     - Uses PyTorch's reduction methods (reduce_tensor, rebuild_tensor, etc.) to handle tensor sharing
     - Supports both CPU tensors (via shared memory) and CUDA tensors (via CUDA IPC)
     - Provides serialization capabilities for cross-process communication
     - Avoids expensive tensor serialization/deserialization/IPC overhead
-    
+
     Architecture:
     - Producer process: Creates tensor -> reduces -> convert handle to dict -> sends via IPC
     - Consumer process: Receives dict -> convert dict to handle -> rebuilds tensor -> gets local view
-    
+
     Note: This is a temporary solution. In the future, if we embrace more Python environment
     or PyTorch orchestration methods (like torch.multiprocessing, Ray, etc.), this intermediate
     layer would not be needed as those frameworks provide native tensor sharing capabilities.
-    
-    Note: Whenever you call reduce_tensor, you must call the corresponding rebuild method at 
-    consumer process(es), otherwise, the producer process cannot release the memory to caching 
+
+    Note: Whenever you call reduce_tensor, you must call the corresponding rebuild method at
+    consumer process(es), otherwise, the producer process cannot release the memory to caching
     allocator as the inner refcount never reaches zero.
     """
+
     def __init__(self, method_key: int, tensor_handle: Dict[str, Any]):
         """Initialize the SharedTensorContainer.
-        
+
         Args:
             method_key: Registry key for the rebuild method (CUDA, CPU, etc.)
             tensor_handle: Tensor handle that can be used to rebuild the tensor on a consumer process
@@ -108,11 +111,11 @@ def __init__(self, method_key: int, tensor_handle: Dict[str, Any]):
     @staticmethod
     def cuda_handle_to_dict(tensor_handle) -> Dict[str, Any]:
         """Convert CUDA tensor handle to serializable dictionary for IPC
-        
+
         This method converts PyTorch's CUDA tensor reduction (cudaIPC) handle into a format that can be
-        safely serialized and transmitted between any two processes. It handles binary data by 
+        safely serialized and transmitted between any two processes. It handles binary data by
         encoding it in base64 to ensure JSON compatibility.
-        
+
         The CUDA handle contains references to GPU memory, CUDA events, and reference counters
         that need to be properly serialized for cross-process sharing via CUDA IPC.
 
@@ -129,19 +132,32 @@ def cuda_handle_to_dict(tensor_handle) -> Dict[str, Any]:
             # Convert tensor info to a basic dict with only serializable values
             serializable_info = {
                 # tensor_info[0] is the type of the tensor, which is "torch.Tensor"
-                "tensor_size": list(tensor_info[1]),
-                "tensor_stride": list(tensor_info[2]),
-                "tensor_offset": tensor_info[3],
-                "dtype": str(tensor_info[5]),
-                "storage_device": tensor_info[6],
-                "storage_handle": base64.b64encode(tensor_info[7]).decode('utf-8'),
-                "storage_size_bytes": tensor_info[8],
-                "storage_offset_bytes": tensor_info[9],
-                "requires_grad": tensor_info[10],
-                "ref_counter_handle": base64.b64encode(tensor_info[11]).decode('utf-8'),
-                "ref_counter_offset": tensor_info[12],
-                "event_handle": base64.b64encode(tensor_info[13]).decode('utf-8'),
-                "event_sync_required": tensor_info[14]
+                "tensor_size":
+                list(tensor_info[1]),
+                "tensor_stride":
+                list(tensor_info[2]),
+                "tensor_offset":
+                tensor_info[3],
+                "dtype":
+                str(tensor_info[5]),
+                "storage_device":
+                tensor_info[6],
+                "storage_handle":
+                base64.b64encode(tensor_info[7]).decode('utf-8'),
+                "storage_size_bytes":
+                tensor_info[8],
+                "storage_offset_bytes":
+                tensor_info[9],
+                "requires_grad":
+                tensor_info[10],
+                "ref_counter_handle":
+                base64.b64encode(tensor_info[11]).decode('utf-8'),
+                "ref_counter_offset":
+                tensor_info[12],
+                "event_handle":
+                base64.b64encode(tensor_info[13]).decode('utf-8'),
+                "event_sync_required":
+                tensor_info[14]
             }
             return serializable_info
         except IndexError as e:
@@ -150,29 +166,37 @@ def cuda_handle_to_dict(tensor_handle) -> Dict[str, Any]:
             raise ValueError(f"Failed to serialize tensor information: {e}")
 
     @staticmethod
-    def cpu_handle_to_dict(meta_data: Dict[str, Any], storage_metadata: Dict[str, Any]) -> Dict[str, Any]:
+    def cpu_handle_to_dict(meta_data: Dict[str, Any],
+                           storage_metadata: Dict[str, Any]) -> Dict[str, Any]:
         """Convert CPU tensor handle to serializable dictionary for IPC
-        
+
         This method converts PyTorch's CPU tensor reduction handle into a format that can be
         safely serialized and transmitted between any two processes. For CPU tensors, we use
         shared memory via file_system sharing strategy as fd strategy is not serializable.
-        
+
         Args:
             meta_data: Tensor metadata (size, stride, offset, etc.)
             storage_metadata: Storage metadata (handle, size, dtype, etc.)
-            
+
         Returns:
             Dictionary containing the serialized CPU tensor information
         """
         try:
             serializable_info = {
-                "tensor_storage_offset": meta_data[0],
-                "tensor_size": list(meta_data[1]),
-                "tensor_stride": list(meta_data[2]),
-                "manager_handle": base64.b64encode(storage_metadata[0]).decode('utf-8'),
-                "storage_handle": base64.b64encode(storage_metadata[1]).decode('utf-8'),
-                "storage_size": storage_metadata[2],
-                "storage_dtype": str(storage_metadata[3])
+                "tensor_storage_offset":
+                meta_data[0],
+                "tensor_size":
+                list(meta_data[1]),
+                "tensor_stride":
+                list(meta_data[2]),
+                "manager_handle":
+                base64.b64encode(storage_metadata[0]).decode('utf-8'),
+                "storage_handle":
+                base64.b64encode(storage_metadata[1]).decode('utf-8'),
+                "storage_size":
+                storage_metadata[2],
+                "storage_dtype":
+                str(storage_metadata[3])
             }
             return serializable_info
         except IndexError as e:
@@ -180,14 +204,13 @@ def cpu_handle_to_dict(meta_data: Dict[str, Any], storage_metadata: Dict[str, An
         except Exception as e:
             raise ValueError(f"Failed to serialize tensor information: {e}")
 
-
     @staticmethod
     def dict_to_cuda_handle(tensor_info: Dict[str, Any]) -> Tuple:
         """Reconstruct CUDA tensor handle from serialized dictionary
-        
+
         This method reconstructs a CUDA tensor handle from a previously serialized dictionary
-        that was received from another process. 
-        
+        that was received from another process.
+
         Args:
             tensor_info: Dictionary containing the serialized CUDA tensor information
                 with the same keys as returned by cuda_handle_to_dict()
@@ -198,24 +221,21 @@ def dict_to_cuda_handle(tensor_info: Dict[str, Any]) -> Tuple:
         try:
             # Decode base64 encoded binary data
             storage_handle = base64.b64decode(tensor_info['storage_handle'])
-            ref_counter_handle = base64.b64decode(tensor_info['ref_counter_handle'])
+            ref_counter_handle = base64.b64decode(
+                tensor_info['ref_counter_handle'])
             event_handle = base64.b64decode(tensor_info['event_handle'])
 
             # Reconstruct the tensor handle
-            tensor_handle = (torch.Tensor,
-                             tuple(tensor_info['tensor_size']),
+            tensor_handle = (torch.Tensor, tuple(tensor_info['tensor_size']),
                              tuple(tensor_info['tensor_stride']),
                              tensor_info['tensor_offset'],
                              torch.storage.TypedStorage,
                              eval(tensor_info['dtype']),
-                             tensor_info['storage_device'],
-                             storage_handle,
+                             tensor_info['storage_device'], storage_handle,
                              tensor_info['storage_size_bytes'],
                              tensor_info['storage_offset_bytes'],
-                             tensor_info['requires_grad'],
-                             ref_counter_handle,
-                             tensor_info['ref_counter_offset'],
-                             event_handle,
+                             tensor_info['requires_grad'], ref_counter_handle,
+                             tensor_info['ref_counter_offset'], event_handle,
                              tensor_info['event_sync_required'])
 
             return tensor_handle
@@ -227,39 +247,36 @@ def dict_to_cuda_handle(tensor_info: Dict[str, Any]) -> Tuple:
     @staticmethod
     def dict_to_cpu_handle(tensor_info: Dict[str, Any]) -> Tuple:
         """Reconstruct CPU tensor handle from serialized dictionary
-        
+
         This method reconstructs a CPU tensor handle from a previously serialized dictionary
         that was received from another process.
-        
+
         The reconstructed handle allows the consumer process to access the same memory
         region that was shared by the producer process, avoiding data copying.
 
         Args:
             tensor_info: Dictionary containing the serialized CPU tensor information
                 with the same keys as returned by cpu_handle_to_dict()
-                
+
         Returns:
             A tuple representing the CPU tensor handle for PyTorch's rebuild_tensor
         """
         try:
             manager_handle = base64.b64decode(tensor_info['manager_handle'])
             storage_handle = base64.b64decode(tensor_info['storage_handle'])
-            storage_metadata = (torch.storage.TypedStorage, 
-                                manager_handle, 
-                                storage_handle, 
-                                tensor_info['storage_size'],
+            storage_metadata = (torch.storage.TypedStorage, manager_handle,
+                                storage_handle, tensor_info['storage_size'],
                                 eval(tensor_info['storage_dtype']))
             storage = rebuild_storage_filename(*storage_metadata)
             if not isinstance(storage, torch.storage.TypedStorage):
-                storage = rebuild_typed_storage(storage, eval(tensor_info['storage_dtype']))
+                storage = rebuild_typed_storage(
+                    storage, eval(tensor_info['storage_dtype']))
 
             meta_data = (tensor_info['tensor_storage_offset'],
                          tuple(tensor_info['tensor_size']),
-                         tuple(tensor_info['tensor_stride']),
-                         False) # requires_grad is always False for cpu tensor
-            tensor_handle = (torch.Tensor,
-                             storage,
-                             meta_data)
+                         tuple(tensor_info['tensor_stride']), False
+                         )  # requires_grad is always False for cpu tensor
+            tensor_handle = (torch.Tensor, storage, meta_data)
             return tensor_handle
         except KeyError as e:
             raise KeyError(f"Missing required tensor information: {e}")
@@ -269,11 +286,11 @@ def dict_to_cpu_handle(tensor_info: Dict[str, Any]) -> Tuple:
     @classmethod
     def from_tensor(cls, tensor: torch.Tensor) -> 'SharedTensorContainer':
         """Create a SharedTensorContainer from a local tensor (Producer side).
-        
+
         This method is called by the producer process to prepare a tensor for sharing
         with other processes. It uses PyTorch's reduction methods to create a handle
         that can be efficiently transmitted and reconstructed by consumer processes.
-        
+
         - CUDA tensors: Uses CUDA IPC for GPU memory sharing
         - CPU tensors: Uses file_system sharing strategy for shared memory
 
@@ -290,66 +307,72 @@ def from_tensor(cls, tensor: torch.Tensor) -> 'SharedTensorContainer':
     @classmethod
     def from_dict(cls, tensor_info: Dict[str, Any]) -> 'SharedTensorContainer':
         """Create a SharedTensorContainer from a serialized dictionary (Consumer side).
-        
+
         This method is called by the consumer process to reconstruct a SharedTensorContainer
         from serialized data received via IPC.
 
         Args:
             tensor_info: Dictionary containing the serialized tensor information
                 received from the producer process via IPC
-                
+
         Returns:
             SharedTensorContainer instance ready for tensor reconstruction
-            
+
         Raises:
             ValueError: If the method_key is not supported
         """
         method_key = tensor_info['method_key']
         if method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CUDA:
-            tensor_handle = SharedTensorContainer.dict_to_cuda_handle(tensor_info)
+            tensor_handle = SharedTensorContainer.dict_to_cuda_handle(
+                tensor_info)
         elif method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU:
-            tensor_handle = SharedTensorContainer.dict_to_cpu_handle(tensor_info)
+            tensor_handle = SharedTensorContainer.dict_to_cpu_handle(
+                tensor_info)
         else:
-            raise ValueError(f"Unsupported shared tensor method key: {method_key}")
+            raise ValueError(
+                f"Unsupported shared tensor method key: {method_key}")
         return cls(method_key, tensor_handle)
 
     def get_local_view(self) -> torch.Tensor:
         """Convert the shared tensor back to a local tensor (Consumer side).
-        
+
         This method is called by the consumer process to obtain the actual tensor
         from the SharedTensorContainer.
 
         Returns:
             The reconstructed tensor
         """
-        rebuild_method = _SharedTensorRebuildMethodRegistry.get_method(self.method_key)
+        rebuild_method = _SharedTensorRebuildMethodRegistry.get_method(
+            self.method_key)
         return rebuild_method(*self.tensor_handle)
 
     def dump_to_dict(self) -> Dict[str, Any]:
         """Convert this container to a dictionary for direct IPC (Producer side).
-        
+
         This method is called by the producer process to serialize the SharedTensorContainer instance
-        into a format that is JSON compatible and can be transmitted via IPC. 
+        into a format that is JSON compatible and can be transmitted via IPC.
 
         Returns:
             Dictionary containing the serialized tensor information
         """
         if self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CUDA:
-            tensor_dict = SharedTensorContainer.cuda_handle_to_dict(self.tensor_handle)
+            tensor_dict = SharedTensorContainer.cuda_handle_to_dict(
+                self.tensor_handle)
         elif self.method_key == _SharedTensorRebuildMethodRegistry.REBUILD_CPU:
             sharing_strategy = get_sharing_strategy()
             # Here we use file_system sharing strategy to make it serializable between two non-python independent processes
             set_sharing_strategy("file_system")
-            storage = self.tensor_handle[1] 
+            storage = self.tensor_handle[1]
             meta_data = self.tensor_handle[2]
             storage_handle = reduce_storage(storage)
             # restore the original sharing strategy
             set_sharing_strategy(sharing_strategy)
             # exclude the first element which is the type of the storage
-            storage_metadata = storage_handle[-1][1:] 
-            tensor_dict = SharedTensorContainer.cpu_handle_to_dict(meta_data, storage_metadata)
+            storage_metadata = storage_handle[-1][1:]
+            tensor_dict = SharedTensorContainer.cpu_handle_to_dict(
+                meta_data, storage_metadata)
         else:
             raise ValueError(f"Unsupported tensor device: {self.method_key}")
-        
+
         tensor_dict["method_key"] = self.method_key
         return tensor_dict
diff --git a/tests/unittest/_torch/multimodal/test_share_tensor.py b/tests/unittest/_torch/multimodal/test_share_tensor.py
index 30fe7a94f827..2165bcdbaf45 100644
--- a/tests/unittest/_torch/multimodal/test_share_tensor.py
+++ b/tests/unittest/_torch/multimodal/test_share_tensor.py
@@ -1,8 +1,10 @@
-import unittest
 import multiprocessing as mp
+import unittest
+
 import torch
 
-from tensorrt_llm._torch.multimodal.mm_utils import SharedTensorContainer, _SharedTensorRebuildMethodRegistry
+from tensorrt_llm._torch.multimodal.mm_utils import (
+    SharedTensorContainer, _SharedTensorRebuildMethodRegistry)
 
 
 class TestShareTensor(unittest.TestCase):
@@ -41,7 +43,8 @@ def test_share_cuda_tensor(self):
         queue = mp.Queue()
 
         # Producer process
-        producer = mp.Process(target=self._producer, args=(queue, self.ref_tensor, "cuda"))
+        producer = mp.Process(target=self._producer,
+                              args=(queue, self.ref_tensor, "cuda"))
         producer.start()
         status, data = queue.get(timeout=100)
         # Verify
@@ -57,7 +60,8 @@ def test_share_cpu_tensor(self):
         queue = mp.Queue()
 
         # Producer process
-        producer = mp.Process(target=self._producer, args=(queue, self.ref_tensor, "cpu"))
+        producer = mp.Process(target=self._producer,
+                              args=(queue, self.ref_tensor, "cpu"))
         producer.start()
         status, data = queue.get(timeout=100)
         # Verify
@@ -71,32 +75,37 @@ def test_share_tensor_different_shapes(self):
         mp.set_start_method('spawn', force=True)
         queue = mp.Queue()
         test_shapes = [
-            (1,),           
-            (2, 3),         
-            (1, 2, 3, 4),   
-            (10,),          
+            (1, ),
+            (2, 3),
+            (1, 2, 3, 4),
+            (10, ),
         ]
 
         for shape in test_shapes:
             with self.subTest(shape=shape):
                 test_tensor = torch.randn(shape)
-                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cpu"))
+                producer = mp.Process(target=self._producer,
+                                      args=(queue, test_tensor, "cpu"))
                 producer.start()
                 status, data = queue.get(timeout=100)
-                
+
                 self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+                reconstructed = SharedTensorContainer.from_dict(
+                    data).get_local_view()
                 self.assertTrue(torch.allclose(reconstructed, test_tensor))
                 producer.join()
-            
+
             with self.subTest(shape=shape):
                 test_tensor = torch.randn(shape)
-                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cuda"))
+                producer = mp.Process(target=self._producer,
+                                      args=(queue, test_tensor, "cuda"))
                 producer.start()
                 status, data = queue.get(timeout=100)
                 self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
-                self.assertTrue(torch.allclose(reconstructed, test_tensor.cuda()))
+                reconstructed = SharedTensorContainer.from_dict(
+                    data).get_local_view()
+                self.assertTrue(
+                    torch.allclose(reconstructed, test_tensor.cuda()))
                 del reconstructed
                 producer.join()
 
@@ -116,30 +125,34 @@ def test_share_tensor_different_dtypes(self):
         for dtype in test_dtypes:
             with self.subTest(dtype=dtype):
                 test_tensor = torch.randn(2, 3).to(dtype)
-                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cpu"))
+                producer = mp.Process(target=self._producer,
+                                      args=(queue, test_tensor, "cpu"))
                 producer.start()
                 status, data = queue.get(timeout=100)
-                
+
                 self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+                reconstructed = SharedTensorContainer.from_dict(
+                    data).get_local_view()
                 self.assertTrue(torch.allclose(reconstructed, test_tensor))
                 self.assertEqual(reconstructed.dtype, test_tensor.dtype)
                 producer.join()
-            
+
             with self.subTest(dtype=dtype):
                 test_tensor = torch.randn(2, 3).to(dtype)
-                producer = mp.Process(target=self._producer, args=(queue, test_tensor, "cuda"))
+                producer = mp.Process(target=self._producer,
+                                      args=(queue, test_tensor, "cuda"))
                 producer.start()
                 status, data = queue.get(timeout=100)
                 self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
-                self.assertTrue(torch.allclose(reconstructed, test_tensor.cuda()))
+                reconstructed = SharedTensorContainer.from_dict(
+                    data).get_local_view()
+                self.assertTrue(
+                    torch.allclose(reconstructed, test_tensor.cuda()))
                 self.assertEqual(reconstructed.dtype, test_tensor.dtype)
                 del reconstructed
                 producer.join()
 
-        
 
 if __name__ == '__main__':
     mp.set_start_method('spawn', force=True)
-    unittest.main() 
\ No newline at end of file
+    unittest.main()

From 180d98fcf03843e865f998ea9430bf286685f3e0 Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Sat, 21 Jun 2025 08:28:08 -0700
Subject: [PATCH 3/7] Remove unused meta tensor

---
 tensorrt_llm/_torch/multimodal/mm_utils.py | 13 +++++--------
 1 file changed, 5 insertions(+), 8 deletions(-)

diff --git a/tensorrt_llm/_torch/multimodal/mm_utils.py b/tensorrt_llm/_torch/multimodal/mm_utils.py
index de2b7defad36..7097d3b4f742 100644
--- a/tensorrt_llm/_torch/multimodal/mm_utils.py
+++ b/tensorrt_llm/_torch/multimodal/mm_utils.py
@@ -4,9 +4,11 @@
 
 import torch
 from torch.multiprocessing import get_sharing_strategy, set_sharing_strategy
-from torch.multiprocessing.reductions import (
-    rebuild_cuda_tensor, rebuild_meta_tensor, rebuild_storage_filename,
-    rebuild_tensor, rebuild_typed_storage, reduce_storage, reduce_tensor)
+from torch.multiprocessing.reductions import (rebuild_cuda_tensor,
+                                              rebuild_storage_filename,
+                                              rebuild_tensor,
+                                              rebuild_typed_storage,
+                                              reduce_storage, reduce_tensor)
 
 logger = logging.getLogger(__name__)
 
@@ -23,7 +25,6 @@ class _SharedTensorRebuildMethodRegistry:
     # Fixed keys for common rebuild methods
     REBUILD_CUDA = 1
     REBUILD_CPU = 2
-    REBUILD_META = 3
 
     _registry: Dict[int, Callable] = {}
 
@@ -33,7 +34,6 @@ def initialize(cls):
         # Register common methods with fixed keys
         cls._registry[cls.REBUILD_CUDA] = rebuild_cuda_tensor
         cls._registry[cls.REBUILD_CPU] = rebuild_tensor
-        cls._registry[cls.REBUILD_META] = rebuild_meta_tensor
 
     @classmethod
     def register(cls, method: Callable) -> int:
@@ -49,9 +49,6 @@ def register(cls, method: Callable) -> int:
             return cls.REBUILD_CUDA
         if method == rebuild_tensor:
             return cls.REBUILD_CPU
-        if method == rebuild_meta_tensor:
-            # TODO: add support for meta tensor if needed
-            return cls.REBUILD_META
         raise NotImplementedError("Other rebuild methods are not supported yet")
 
     @classmethod

From 5fb88fbd9ab6da2b2f9dd9c793e3578e3b511b4e Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Sat, 21 Jun 2025 18:37:02 -0700
Subject: [PATCH 4/7] Avoid using eval()

---
 tensorrt_llm/_torch/multimodal/__init__.py |  8 +++++++
 tensorrt_llm/_torch/multimodal/mm_utils.py | 25 +++++++++++++++++++---
 2 files changed, 30 insertions(+), 3 deletions(-)
 create mode 100644 tensorrt_llm/_torch/multimodal/__init__.py

diff --git a/tensorrt_llm/_torch/multimodal/__init__.py b/tensorrt_llm/_torch/multimodal/__init__.py
new file mode 100644
index 000000000000..3c51243e1d55
--- /dev/null
+++ b/tensorrt_llm/_torch/multimodal/__init__.py
@@ -0,0 +1,8 @@
+from .mm_utils import _SharedTensorRebuildMethodRegistry, SharedTensorContainer
+
+# Initialize the registry when the package is imported
+_SharedTensorRebuildMethodRegistry.initialize()
+
+__all__ = [
+    'SharedTensorContainer',
+]
diff --git a/tensorrt_llm/_torch/multimodal/mm_utils.py b/tensorrt_llm/_torch/multimodal/mm_utils.py
index 7097d3b4f742..5d25a7766813 100644
--- a/tensorrt_llm/_torch/multimodal/mm_utils.py
+++ b/tensorrt_llm/_torch/multimodal/mm_utils.py
@@ -12,6 +12,25 @@
 
 logger = logging.getLogger(__name__)
 
+DTYPE_MAPPING = {
+    'torch.float32': torch.float32,
+    'torch.float64': torch.float64,
+    'torch.int32': torch.int32,
+    'torch.int64': torch.int64,
+    'torch.bool': torch.bool,
+    'torch.uint8': torch.uint8,
+    'torch.int8': torch.int8,
+    'torch.int16': torch.int16,
+    'torch.float16': torch.float16,
+    'torch.bfloat16': torch.bfloat16,
+}
+
+def str_to_torch_dtype(dtype_str: str) -> torch.dtype:
+    """Convert dtype string to torch dtype
+    """
+    if dtype_str not in DTYPE_MAPPING:
+        raise ValueError(f"Unsupported dtype: {dtype_str}. Supported dtypes are: {list(DTYPE_MAPPING.keys())}")
+    return DTYPE_MAPPING[dtype_str]
 
 class _SharedTensorRebuildMethodRegistry:
     """Registry for tensor rebuild methods with fixed keys for common methods.
@@ -227,7 +246,7 @@ def dict_to_cuda_handle(tensor_info: Dict[str, Any]) -> Tuple:
                              tuple(tensor_info['tensor_stride']),
                              tensor_info['tensor_offset'],
                              torch.storage.TypedStorage,
-                             eval(tensor_info['dtype']),
+                             str_to_torch_dtype(tensor_info['dtype']),
                              tensor_info['storage_device'], storage_handle,
                              tensor_info['storage_size_bytes'],
                              tensor_info['storage_offset_bytes'],
@@ -263,11 +282,11 @@ def dict_to_cpu_handle(tensor_info: Dict[str, Any]) -> Tuple:
             storage_handle = base64.b64decode(tensor_info['storage_handle'])
             storage_metadata = (torch.storage.TypedStorage, manager_handle,
                                 storage_handle, tensor_info['storage_size'],
-                                eval(tensor_info['storage_dtype']))
+                                str_to_torch_dtype(tensor_info['storage_dtype']))
             storage = rebuild_storage_filename(*storage_metadata)
             if not isinstance(storage, torch.storage.TypedStorage):
                 storage = rebuild_typed_storage(
-                    storage, eval(tensor_info['storage_dtype']))
+                    storage, str_to_torch_dtype(tensor_info['storage_dtype']))
 
             meta_data = (tensor_info['tensor_storage_offset'],
                          tuple(tensor_info['tensor_size']),

From 2f1238e899d70a7071b6a601680a13b778cbb5fa Mon Sep 17 00:00:00 2001
From: "Chang Liu (Enterprise Products)"
 <9713593+chang-l@users.noreply.github.com>
Date: Sat, 21 Jun 2025 18:43:01 -0700
Subject: [PATCH 5/7] Formatting

---
 tensorrt_llm/_torch/multimodal/__init__.py | 2 +-
 tensorrt_llm/_torch/multimodal/mm_utils.py | 9 +++++++--
 2 files changed, 8 insertions(+), 3 deletions(-)

diff --git a/tensorrt_llm/_torch/multimodal/__init__.py b/tensorrt_llm/_torch/multimodal/__init__.py
index 3c51243e1d55..d388edf89506 100644
--- a/tensorrt_llm/_torch/multimodal/__init__.py
+++ b/tensorrt_llm/_torch/multimodal/__init__.py
@@ -1,4 +1,4 @@
-from .mm_utils import _SharedTensorRebuildMethodRegistry, SharedTensorContainer
+from .mm_utils import SharedTensorContainer, _SharedTensorRebuildMethodRegistry
 
 # Initialize the registry when the package is imported
 _SharedTensorRebuildMethodRegistry.initialize()
diff --git a/tensorrt_llm/_torch/multimodal/mm_utils.py b/tensorrt_llm/_torch/multimodal/mm_utils.py
index 5d25a7766813..89216ff03d52 100644
--- a/tensorrt_llm/_torch/multimodal/mm_utils.py
+++ b/tensorrt_llm/_torch/multimodal/mm_utils.py
@@ -25,13 +25,17 @@
     'torch.bfloat16': torch.bfloat16,
 }
 
+
 def str_to_torch_dtype(dtype_str: str) -> torch.dtype:
     """Convert dtype string to torch dtype
     """
     if dtype_str not in DTYPE_MAPPING:
-        raise ValueError(f"Unsupported dtype: {dtype_str}. Supported dtypes are: {list(DTYPE_MAPPING.keys())}")
+        raise ValueError(
+            f"Unsupported dtype: {dtype_str}. Supported dtypes are: {list(DTYPE_MAPPING.keys())}"
+        )
     return DTYPE_MAPPING[dtype_str]
 
+
 class _SharedTensorRebuildMethodRegistry:
     """Registry for tensor rebuild methods with fixed keys for common methods.
 
@@ -282,7 +286,8 @@ def dict_to_cpu_handle(tensor_info: Dict[str, Any]) -> Tuple:
             storage_handle = base64.b64decode(tensor_info['storage_handle'])
             storage_metadata = (torch.storage.TypedStorage, manager_handle,
                                 storage_handle, tensor_info['storage_size'],
-                                str_to_torch_dtype(tensor_info['storage_dtype']))
+                                str_to_torch_dtype(
+                                    tensor_info['storage_dtype']))
             storage = rebuild_storage_filename(*storage_metadata)
             if not isinstance(storage, torch.storage.TypedStorage):
                 storage = rebuild_typed_storage(

From ef2cc2ab57d32bbc2d6480740fd2d16f42e277f3 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Tue, 8 Jul 2025 15:01:35 -0700
Subject: [PATCH 6/7] Address comments and add more tests

---
 .../{multimodal => shared_tensor}/__init__.py |   3 +-
 .../shared_tensor.py}                         |   5 +
 .../{multimodal => }/test_share_tensor.py     | 107 ++++++++++++++++--
 3 files changed, 107 insertions(+), 8 deletions(-)
 rename tensorrt_llm/_torch/{multimodal => shared_tensor}/__init__.py (56%)
 rename tensorrt_llm/_torch/{multimodal/mm_utils.py => shared_tensor/shared_tensor.py} (97%)
 rename tests/unittest/_torch/{multimodal => }/test_share_tensor.py (60%)

diff --git a/tensorrt_llm/_torch/multimodal/__init__.py b/tensorrt_llm/_torch/shared_tensor/__init__.py
similarity index 56%
rename from tensorrt_llm/_torch/multimodal/__init__.py
rename to tensorrt_llm/_torch/shared_tensor/__init__.py
index d388edf89506..d7548f9391e0 100644
--- a/tensorrt_llm/_torch/multimodal/__init__.py
+++ b/tensorrt_llm/_torch/shared_tensor/__init__.py
@@ -1,4 +1,5 @@
-from .mm_utils import SharedTensorContainer, _SharedTensorRebuildMethodRegistry
+from .shared_tensor import (SharedTensorContainer,
+                            _SharedTensorRebuildMethodRegistry)
 
 # Initialize the registry when the package is imported
 _SharedTensorRebuildMethodRegistry.initialize()
diff --git a/tensorrt_llm/_torch/multimodal/mm_utils.py b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py
similarity index 97%
rename from tensorrt_llm/_torch/multimodal/mm_utils.py
rename to tensorrt_llm/_torch/shared_tensor/shared_tensor.py
index 89216ff03d52..b210f1b24632 100644
--- a/tensorrt_llm/_torch/multimodal/mm_utils.py
+++ b/tensorrt_llm/_torch/shared_tensor/shared_tensor.py
@@ -116,6 +116,11 @@ class SharedTensorContainer:
     Note: Whenever you call reduce_tensor, you must call the corresponding rebuild method at
     consumer process(es), otherwise, the producer process cannot release the memory to caching
     allocator as the inner refcount never reaches zero.
+
+    Note: This module can also be extended to transfer CUDA tensors between different GPUs managed by different processes
+    using CE (Copy Engine) or torch.to() to initiate direct P2P transfers. This requires CUDA P2P support, i.e.,
+    torch.cuda.can_device_access_peer(src_device, dst_device) must return True for the source and destination devices.
+
     """
 
     def __init__(self, method_key: int, tensor_handle: Dict[str, Any]):
diff --git a/tests/unittest/_torch/multimodal/test_share_tensor.py b/tests/unittest/_torch/test_share_tensor.py
similarity index 60%
rename from tests/unittest/_torch/multimodal/test_share_tensor.py
rename to tests/unittest/_torch/test_share_tensor.py
index 2165bcdbaf45..09180347840d 100644
--- a/tests/unittest/_torch/multimodal/test_share_tensor.py
+++ b/tests/unittest/_torch/test_share_tensor.py
@@ -3,18 +3,12 @@
 
 import torch
 
-from tensorrt_llm._torch.multimodal.mm_utils import (
-    SharedTensorContainer, _SharedTensorRebuildMethodRegistry)
+from tensorrt_llm._torch.shared_tensor import SharedTensorContainer
 
 
 class TestShareTensor(unittest.TestCase):
     """Test cases for sharing tensors between processes."""
 
-    @classmethod
-    def setUpClass(cls):
-        """Initialize the registry before running tests."""
-        _SharedTensorRebuildMethodRegistry.initialize()
-
     def setUp(self):
         """Set up test fixtures."""
         self.ref_tensor = torch.randn(3, 4, 5)
@@ -33,6 +27,11 @@ def _producer(q, tensor, device=None):
                     tensor = tensor.cpu()
             container = SharedTensorContainer.from_tensor(tensor)
             q.put(('success', container.dump_to_dict()))
+
+            # Wait for consumer to signal it's done
+            # This keeps the producer alive until ownership is transferred to consumer
+            q.get()
+
         except Exception as e:
             q.put(('error', str(e)))
 
@@ -50,6 +49,9 @@ def test_share_cuda_tensor(self):
         # Verify
         self.assertEqual(status, 'success')
         reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+        queue.put(
+            'done'
+        )  # producer can be released as early as here as ownership is transferred to consumer
         self.assertTrue(torch.allclose(reconstructed.cpu(), self.ref_tensor))
         del reconstructed
         producer.join()
@@ -67,6 +69,9 @@ def test_share_cpu_tensor(self):
         # Verify
         self.assertEqual(status, 'success')
         reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
+        queue.put(
+            'done'
+        )  # producer can be released as early as here as ownership is transferred to consumer
         self.assertTrue(torch.allclose(reconstructed, self.ref_tensor))
         producer.join()
 
@@ -93,6 +98,7 @@ def test_share_tensor_different_shapes(self):
                 reconstructed = SharedTensorContainer.from_dict(
                     data).get_local_view()
                 self.assertTrue(torch.allclose(reconstructed, test_tensor))
+                queue.put('done')
                 producer.join()
 
             with self.subTest(shape=shape):
@@ -107,6 +113,7 @@ def test_share_tensor_different_shapes(self):
                 self.assertTrue(
                     torch.allclose(reconstructed, test_tensor.cuda()))
                 del reconstructed
+                queue.put('done')
                 producer.join()
 
     def test_share_tensor_different_dtypes(self):
@@ -135,6 +142,7 @@ def test_share_tensor_different_dtypes(self):
                     data).get_local_view()
                 self.assertTrue(torch.allclose(reconstructed, test_tensor))
                 self.assertEqual(reconstructed.dtype, test_tensor.dtype)
+                queue.put('done')
                 producer.join()
 
             with self.subTest(dtype=dtype):
@@ -150,8 +158,93 @@ def test_share_tensor_different_dtypes(self):
                     torch.allclose(reconstructed, test_tensor.cuda()))
                 self.assertEqual(reconstructed.dtype, test_tensor.dtype)
                 del reconstructed
+                queue.put('done')
                 producer.join()
 
+    @staticmethod
+    def _stand_by_producer(conn):
+        """Long-lived producer that creates new tensors on demand."""
+        try:
+            while True:
+                msg = conn.recv()
+                if msg == "get":
+                    # Create a new tensor each time
+                    tensor = torch.randn(100, 100, 100).cuda()  # ~4MB tensor
+                    container = SharedTensorContainer.from_tensor(tensor)
+                    serialized_data = container.dump_to_dict()
+                    memory_usage = torch.cuda.memory_allocated() / (1024 * 1024)
+                    conn.send(('success', serialized_data, memory_usage))
+                elif msg == "exit":
+                    break
+                else:
+                    print(f"Unknown command: {msg}")
+        except Exception as e:
+            conn.send(('error', str(e)))
+        finally:
+            conn.close()
+
+    @unittest.skipIf(not torch.cuda.is_available(), "CUDA not available")
+    def test_memory_leak_repeated_producer(self):
+        """Test to check no memory leak when producer creates new tensors repeatedly.
+
+        This test keeps the producer alive and requests multiple tensors.
+        Each iteration, the producer creates a new tensor and shares it.
+        If the consumer properly rebuild and cleanup, GPU memory usage will likely be stable.
+        """
+        import gc
+
+        import numpy as np
+
+        mp.set_start_method('spawn', force=True)
+
+        # Reset GPU state before test
+        torch.cuda.empty_cache()
+        torch.cuda.synchronize()
+
+        # Record initial memory state
+        initial_memory = torch.cuda.memory_allocated() / (1024 * 1024)
+
+        parent_conn, child_conn = mp.Pipe()
+        producer = mp.Process(target=self._stand_by_producer,
+                              args=(child_conn, ))
+        producer.start()
+
+        memory_measurements = []
+        try:
+            for i in range(10):
+                parent_conn.send("get")
+                status, data, memory_usage = parent_conn.recv()
+                memory_measurements.append(memory_usage)
+                self.assertEqual(status, 'success')
+
+                container = SharedTensorContainer.from_dict(
+                    data).get_local_view()
+                del container
+                gc.collect()
+                torch.cuda.ipc_collect()
+                torch.cuda.empty_cache()
+                torch.cuda.synchronize()
+
+            relative_measurements = [
+                m - initial_memory for m in memory_measurements
+            ]
+
+            warmup_iterations = 4
+            stable_measurements = relative_measurements[warmup_iterations:]
+
+            x = np.arange(len(stable_measurements))
+            slope, _ = np.polyfit(x, stable_measurements, 1)
+
+            self.assertLess(
+                abs(slope), 0.2,
+                f"Memory leak detected! Relative slope: {slope:.3f} MB/iteration. "
+                f"Relative measurements: {relative_measurements}")
+
+        finally:
+            parent_conn.send("exit")
+            producer.join()
+            parent_conn.close()
+
 
 if __name__ == '__main__':
     mp.set_start_method('spawn', force=True)

From 3103e663a118f100bbc96e0faf97970702cf73d5 Mon Sep 17 00:00:00 2001
From: Chang Liu <9713593+chang-l@users.noreply.github.com>
Date: Wed, 9 Jul 2025 10:04:08 -0700
Subject: [PATCH 7/7] Cleanup queue to prevent thread leak in tests

---
 tests/unittest/_torch/test_share_tensor.py | 197 ++++++++++++---------
 1 file changed, 110 insertions(+), 87 deletions(-)

diff --git a/tests/unittest/_torch/test_share_tensor.py b/tests/unittest/_torch/test_share_tensor.py
index 09180347840d..548148255ebc 100644
--- a/tests/unittest/_torch/test_share_tensor.py
+++ b/tests/unittest/_torch/test_share_tensor.py
@@ -41,39 +41,52 @@ def test_share_cuda_tensor(self):
         mp.set_start_method('spawn', force=True)
         queue = mp.Queue()
 
-        # Producer process
-        producer = mp.Process(target=self._producer,
-                              args=(queue, self.ref_tensor, "cuda"))
-        producer.start()
-        status, data = queue.get(timeout=100)
-        # Verify
-        self.assertEqual(status, 'success')
-        reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
-        queue.put(
-            'done'
-        )  # producer can be released as early as here as ownership is transferred to consumer
-        self.assertTrue(torch.allclose(reconstructed.cpu(), self.ref_tensor))
-        del reconstructed
-        producer.join()
+        try:
+            # Producer process
+            producer = mp.Process(target=self._producer,
+                                  args=(queue, self.ref_tensor, "cuda"))
+            producer.start()
+            status, data = queue.get(timeout=100)
+            # Verify
+            self.assertEqual(status, 'success')
+            reconstructed = SharedTensorContainer.from_dict(
+                data).get_local_view()
+            queue.put(
+                'done'
+            )  # producer can be released as early as here as ownership is transferred to consumer
+            self.assertTrue(torch.allclose(reconstructed.cpu(),
+                                           self.ref_tensor))
+            del reconstructed
+            producer.join()
+        finally:
+            # Explicit cleanup to prevent QueueFeederThread leak
+            queue.close()
+            queue.join_thread()
 
     def test_share_cpu_tensor(self):
         """Test CPU tensor sharing between processes."""
         mp.set_start_method('spawn', force=True)
         queue = mp.Queue()
 
-        # Producer process
-        producer = mp.Process(target=self._producer,
-                              args=(queue, self.ref_tensor, "cpu"))
-        producer.start()
-        status, data = queue.get(timeout=100)
-        # Verify
-        self.assertEqual(status, 'success')
-        reconstructed = SharedTensorContainer.from_dict(data).get_local_view()
-        queue.put(
-            'done'
-        )  # producer can be released as early as here as ownership is transferred to consumer
-        self.assertTrue(torch.allclose(reconstructed, self.ref_tensor))
-        producer.join()
+        try:
+            # Producer process
+            producer = mp.Process(target=self._producer,
+                                  args=(queue, self.ref_tensor, "cpu"))
+            producer.start()
+            status, data = queue.get(timeout=100)
+            # Verify
+            self.assertEqual(status, 'success')
+            reconstructed = SharedTensorContainer.from_dict(
+                data).get_local_view()
+            queue.put(
+                'done'
+            )  # producer can be released as early as here as ownership is transferred to consumer
+            self.assertTrue(torch.allclose(reconstructed, self.ref_tensor))
+            producer.join()
+        finally:
+            # Explicit cleanup to prevent QueueFeederThread leak
+            queue.close()
+            queue.join_thread()
 
     def test_share_tensor_different_shapes(self):
         """Test CPU tensor sharing with different tensor shapes."""
@@ -86,35 +99,40 @@ def test_share_tensor_different_shapes(self):
             (10, ),
         ]
 
-        for shape in test_shapes:
-            with self.subTest(shape=shape):
-                test_tensor = torch.randn(shape)
-                producer = mp.Process(target=self._producer,
-                                      args=(queue, test_tensor, "cpu"))
-                producer.start()
-                status, data = queue.get(timeout=100)
-
-                self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(
-                    data).get_local_view()
-                self.assertTrue(torch.allclose(reconstructed, test_tensor))
-                queue.put('done')
-                producer.join()
-
-            with self.subTest(shape=shape):
-                test_tensor = torch.randn(shape)
-                producer = mp.Process(target=self._producer,
-                                      args=(queue, test_tensor, "cuda"))
-                producer.start()
-                status, data = queue.get(timeout=100)
-                self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(
-                    data).get_local_view()
-                self.assertTrue(
-                    torch.allclose(reconstructed, test_tensor.cuda()))
-                del reconstructed
-                queue.put('done')
-                producer.join()
+        try:
+            for shape in test_shapes:
+                with self.subTest(shape=shape):
+                    test_tensor = torch.randn(shape)
+                    producer = mp.Process(target=self._producer,
+                                          args=(queue, test_tensor, "cpu"))
+                    producer.start()
+                    status, data = queue.get(timeout=100)
+
+                    self.assertEqual(status, 'success')
+                    reconstructed = SharedTensorContainer.from_dict(
+                        data).get_local_view()
+                    self.assertTrue(torch.allclose(reconstructed, test_tensor))
+                    queue.put('done')
+                    producer.join()
+
+                with self.subTest(shape=shape):
+                    test_tensor = torch.randn(shape)
+                    producer = mp.Process(target=self._producer,
+                                          args=(queue, test_tensor, "cuda"))
+                    producer.start()
+                    status, data = queue.get(timeout=100)
+                    self.assertEqual(status, 'success')
+                    reconstructed = SharedTensorContainer.from_dict(
+                        data).get_local_view()
+                    self.assertTrue(
+                        torch.allclose(reconstructed, test_tensor.cuda()))
+                    del reconstructed
+                    queue.put('done')
+                    producer.join()
+        finally:
+            # Explicit cleanup to prevent QueueFeederThread leak
+            queue.close()
+            queue.join_thread()
 
     def test_share_tensor_different_dtypes(self):
         """Test CPU tensor sharing with different data types."""
@@ -129,37 +147,42 @@ def test_share_tensor_different_dtypes(self):
             torch.int64,
         ]
 
-        for dtype in test_dtypes:
-            with self.subTest(dtype=dtype):
-                test_tensor = torch.randn(2, 3).to(dtype)
-                producer = mp.Process(target=self._producer,
-                                      args=(queue, test_tensor, "cpu"))
-                producer.start()
-                status, data = queue.get(timeout=100)
-
-                self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(
-                    data).get_local_view()
-                self.assertTrue(torch.allclose(reconstructed, test_tensor))
-                self.assertEqual(reconstructed.dtype, test_tensor.dtype)
-                queue.put('done')
-                producer.join()
-
-            with self.subTest(dtype=dtype):
-                test_tensor = torch.randn(2, 3).to(dtype)
-                producer = mp.Process(target=self._producer,
-                                      args=(queue, test_tensor, "cuda"))
-                producer.start()
-                status, data = queue.get(timeout=100)
-                self.assertEqual(status, 'success')
-                reconstructed = SharedTensorContainer.from_dict(
-                    data).get_local_view()
-                self.assertTrue(
-                    torch.allclose(reconstructed, test_tensor.cuda()))
-                self.assertEqual(reconstructed.dtype, test_tensor.dtype)
-                del reconstructed
-                queue.put('done')
-                producer.join()
+        try:
+            for dtype in test_dtypes:
+                with self.subTest(dtype=dtype):
+                    test_tensor = torch.randn(2, 3).to(dtype)
+                    producer = mp.Process(target=self._producer,
+                                          args=(queue, test_tensor, "cpu"))
+                    producer.start()
+                    status, data = queue.get(timeout=100)
+
+                    self.assertEqual(status, 'success')
+                    reconstructed = SharedTensorContainer.from_dict(
+                        data).get_local_view()
+                    self.assertTrue(torch.allclose(reconstructed, test_tensor))
+                    self.assertEqual(reconstructed.dtype, test_tensor.dtype)
+                    queue.put('done')
+                    producer.join()
+
+                with self.subTest(dtype=dtype):
+                    test_tensor = torch.randn(2, 3).to(dtype)
+                    producer = mp.Process(target=self._producer,
+                                          args=(queue, test_tensor, "cuda"))
+                    producer.start()
+                    status, data = queue.get(timeout=100)
+                    self.assertEqual(status, 'success')
+                    reconstructed = SharedTensorContainer.from_dict(
+                        data).get_local_view()
+                    self.assertTrue(
+                        torch.allclose(reconstructed, test_tensor.cuda()))
+                    self.assertEqual(reconstructed.dtype, test_tensor.dtype)
+                    del reconstructed
+                    queue.put('done')
+                    producer.join()
+        finally:
+            # Explicit cleanup to prevent QueueFeederThread leak
+            queue.close()
+            queue.join_thread()
 
     @staticmethod
     def _stand_by_producer(conn):