[Bugfix] Fix PyTorch H2D input lifetime across CUDA streams

lmdeploy/pytorch/engine/logits_process.py

@@ -285,6 +285,13 @@ def to_device(self, device: str, non_blocking: bool = False):
 
         return SamplingInputs(**out_dict)
 
+    def record_stream(self, stream: torch.cuda.Stream) -> None:
+        """Record forward-stream use of tensor fields."""
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if isinstance(value, torch.Tensor) and value.is_cuda:
+                value.record_stream(stream)
+
     def get_delta(self) -> SamplingInputsDelta:
         """Get delta."""
         delta = SamplingInputsDelta()

lmdeploy/pytorch/engine/model_agent/agent.py

@@ -2,6 +2,7 @@
 import asyncio
 import time
 from collections import deque
+from collections.abc import Mapping
 from contextlib import contextmanager
 from dataclasses import dataclass, field, fields
 from multiprocessing.reduction import ForkingPickler
@@ -46,6 +47,7 @@
 logger = get_logger('lmdeploy')
 
 _H2D_TRANSFER_KEY = '_h2d_transfer'
+_H2D_INPUT_KEYS = ('inputs', 'delta', 'sampling_inputs', 'stopping_criteria', 'extra_inputs')
 
 
 @dataclass
@@ -236,6 +238,23 @@ def _try_to_cuda(val, non_blocking: bool = False):
         raise RuntimeError(f'Can not cast {type(val)} to cuda.')
 
 
+def _record_forward_input_stream(forward_inputs: Mapping[str, Any], stream: torch.cuda.Stream) -> None:
+    """Register the forward stream with each H2D payload owner."""
+    for key in _H2D_INPUT_KEYS:
+        value = forward_inputs.get(key)
+        if value is None:
+            continue
+        if isinstance(value, torch.Tensor):
+            if value.is_cuda:
+                value.record_stream(stream)
+            continue
+
+        record_stream = getattr(value, 'record_stream', None)
+        if not callable(record_stream):
+            raise TypeError(f'H2D input {key!r} must be a tensor or implement record_stream().')
+        record_stream(stream)
+
+
 SwapMap = dict[int, int]
 
 
@@ -938,6 +957,7 @@ async def _async_loop_background(self, forward_event: asyncio.Event = None):
                 if h2d_transfer is not None:
                     self._keep_h2d_transfer(h2d_transfer)
                     self.stream.wait_event(h2d_transfer.event)
+                    _record_forward_input_stream(forward_inputs, self.stream)
 
                 await self._async_step(**forward_inputs, )
                 if forward_event is not None:
@@ -970,18 +990,20 @@ async def _async_loop_inputs_preprocess(self, forward_event: asyncio.Event = Non
         falling back to dummy inputs.
         """
         non_blocking = True
-        keys = ['inputs', 'delta', 'sampling_inputs', 'stopping_criteria', 'extra_inputs']
         while True:
             forward_inputs = await self._pre_in_que.get()
             forward_inputs_cuda = {}
             forward_inputs_cuda.update(forward_inputs)
-            h2d_refs = {k: forward_inputs_cuda[k] for k in keys if forward_inputs_cuda.get(k) is not None}
+            h2d_refs = {
+                key: forward_inputs_cuda[key]
+                for key in _H2D_INPUT_KEYS if forward_inputs_cuda.get(key) is not None
+            }
             logger.debug('preprocessing forward inputs.')
             with torch.cuda.stream(self.out_stream), torch.inference_mode(), record_function('inputs_H2D'):
-                for k in keys:
-                    if k not in forward_inputs_cuda:
+                for key in _H2D_INPUT_KEYS:
+                    if key not in forward_inputs_cuda:
                         continue
-                    forward_inputs_cuda[k] = _try_to_cuda(forward_inputs_cuda[k], non_blocking=non_blocking)
+                    forward_inputs_cuda[key] = _try_to_cuda(forward_inputs_cuda[key], non_blocking=non_blocking)
                 h2d_event = torch.cuda.Event()
                 h2d_event.record()
                 forward_inputs_cuda[_H2D_TRANSFER_KEY] = _H2DTransfer(h2d_event, h2d_refs)

lmdeploy/pytorch/model_inputs.py

@@ -101,6 +101,29 @@ def to_device(self, device: str, non_blocking: bool = False):
 
         return VisionModelInputs(**out_dict)
 
+    def record_stream(self, stream: torch.cuda.Stream) -> None:
+        """Record forward-stream use of vision tensor fields."""
+        for f in fields(self):
+            key = f.name
+            value = getattr(self, key)
+            if isinstance(value, torch.Tensor):
+                if value.is_cuda:
+                    value.record_stream(stream)
+            elif key == 'input_embedding_ranges':
+                for tensor in value or ():
+                    if tensor.is_cuda:
+                        tensor.record_stream(stream)
+            elif key == 'input_embeddings':
+                for tensors in value or ():
+                    for tensor in tensors:
+                        if tensor.is_cuda:
+                            tensor.record_stream(stream)
+            elif key == 'input_multimodals':
+                for multimodals in value or ():
+                    for items in multimodals.values():
+                        for item in items:
+                            item.record_stream(stream)
+
     def get_inputs(self, history_lengths: torch.Tensor, seq_lengths: torch.Tensor):
         """Get vision embedding inputs."""
         input_embeddings = None
@@ -172,6 +195,13 @@ def to_device(self, device: str, non_blocking: bool = False):
 
         return ModelInputsDelta(**out_dict)
 
+    def record_stream(self, stream: torch.cuda.Stream) -> None:
+        """Record forward-stream use of tensor fields."""
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if isinstance(value, torch.Tensor) and value.is_cuda:
+                value.record_stream(stream)
+
     def log_info(self):
         """Get log info."""
         ret = (f'num_tokens={self.indices.numel()}, batch_size={self.indices.numel()}'
@@ -258,6 +288,16 @@ def to_device(self, device: str, non_blocking: bool = False):
 
         return ModelInputs(**out_dict)
 
+    def record_stream(self, stream: torch.cuda.Stream) -> None:
+        """Record forward-stream use of model tensor fields."""
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if isinstance(value, torch.Tensor):
+                if value.is_cuda:
+                    value.record_stream(stream)
+            elif isinstance(value, VisionModelInputs):
+                value.record_stream(stream)
+
     def build_dp_meta(self, num_tokens: list[int]):
         """Build dp meta."""
         self.dp_meta = DPMeta.build(self.input_ids.numel(), num_tokens)

lmdeploy/pytorch/multimodal/data_type.py

@@ -99,6 +99,20 @@ def to_device(self, device: str, non_blocking: bool = False):
         out_dict['meta'] = new_meta
         return MultiModalData(**out_dict)
 
+    def record_stream(self, stream: torch.cuda.Stream) -> None:
+        """Record forward-stream use of multimodal tensor fields."""
+        tensors = [self.data] if isinstance(self.data, Tensor) else self.data
+        for tensor in tensors:
+            if tensor.is_cuda:
+                tensor.record_stream(stream)
+
+        for value in (self.meta or {}).values():
+            if isinstance(value, Tensor):
+                if value.is_cuda:
+                    value.record_stream(stream)
+            elif hasattr(value, 'record_stream'):
+                value.record_stream(stream)
+
 
 MultiModalInputs = dict[str, list[MultiModalData]]
 

lmdeploy/pytorch/strategies/base/model_agent.py

@@ -35,6 +35,13 @@ def to_device(self, device: str, non_blocking: bool = False):
         """To device."""
         return to_device(self, device, non_blocking)
 
+    def record_stream(self, stream: torch.cuda.Stream) -> None:
+        """Record forward-stream use of tensor fields."""
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if isinstance(value, torch.Tensor) and value.is_cuda:
+                value.record_stream(stream)
+
     def broadcast(self, src: int, group, async_op=False):
         """Broadcast extra inputs."""
         pass
@@ -111,6 +118,13 @@ def to_device(self, device: str, non_blocking: bool = False):
         """To device."""
         return to_device(self, device, non_blocking)
 
+    def record_stream(self, stream: torch.cuda.Stream) -> None:
+        """Record forward-stream use of tensor fields."""
+        for f in fields(self):
+            value = getattr(self, f.name)
+            if isinstance(value, torch.Tensor) and value.is_cuda:
+                value.record_stream(stream)
+
 
 class ModelAgentStrategy(ABC):
     """Base class for model agent strategies."""

tests/pytorch/engine/test_logits_process.py

@@ -11,6 +11,35 @@
 # yapf: enable
 
 
+def test_sampling_inputs_record_stream_records_only_tensor_fields():
+    from lmdeploy.pytorch.engine.logits_process import SamplingInputs
+
+    recorded = []
+
+    class _CudaTensor(torch.Tensor):
+
+        @staticmethod
+        def __new__(cls):
+            return torch.Tensor._make_subclass(cls, torch.empty(1), False)
+
+        @property
+        def is_cuda(self):
+            return True
+
+        def record_stream(self, stream):
+            recorded.append((id(self), stream))
+
+    stream = object()
+    temperature = _CudaTensor()
+    nested_session_tensor = _CudaTensor()
+    inputs = SamplingInputs(temperature=temperature,
+                            session_ctx=[{'persistent': nested_session_tensor}])
+
+    inputs.record_stream(stream)
+
+    assert recorded == [(id(temperature), stream)]
+
+
 def test_process_temperature():
     from lmdeploy.pytorch.engine.logits_process import _process_temperature_