[Bugfix] Fix PyTorch H2D input lifetime across CUDA streams

[grimoire]

Motivation

The PyTorch model agent transfers forward inputs on the H2D stream and
consumes the resulting tensors on the forward stream after waiting for a
CUDA event.

The event orders execution, but does not register the forward stream with
PyTorch's caching allocator. After Python releases the tensor references,
the allocations could therefore be reused on the H2D stream while
forward-stream kernels were still consuming them.

_H2DTransfer.refs protects CPU source objects until the non-blocking H2D
copies complete, but does not protect the lifetime of the resulting device
allocations.

Modification

• Use a shared input-key tuple for H2D transfer, source retention, and
  forward-stream lifetime registration.
• Record the forward stream on every transferred payload after waiting for
  the H2D event and before model execution.
• Add owner-local record_stream() methods for model, delta, vision,
  multimodal, sampling, stopping, and strategy extra inputs.
• Fail early when a non-tensor H2D payload does not implement the lifetime
  protocol.

This preserves the existing asynchronous H2D/forward overlap and does not
introduce a stream synchronization.

[Copilot]

Pull request overview

Fixes a CUDA caching-allocator lifetime hazard in the PyTorch model agent by registering the forward CUDA stream against H2D-transferred device allocations after waiting on the H2D event, preventing premature reuse across streams.

Changes:

• Introduces _record_forward_input_stream() and integrates it into the background loop after wait_event() and before model execution.
• Adds record_stream() implementations across model/vision/delta/sampling inputs and strategy extra inputs/stopping criteria.
• Adds unit tests to validate stream-recording behavior and a CUDA allocator regression test.

Reviewed changes

Copilot reviewed 8 out of 8 changed files in this pull request and generated 1 comment.

Show a summary per file

File	Description
tests/pytorch/test_model_inputs.py	Adds tests verifying ModelInputs/VisionModelInputs and ModelInputsDelta stream recording.
tests/pytorch/engine/test_model_agent.py	Adds tests for _record_forward_input_stream() protocol behavior and an allocator-lifetime regression test.
tests/pytorch/engine/test_logits_process.py	Adds a test for SamplingInputs.record_stream() behavior.
lmdeploy/pytorch/strategies/base/model_agent.py	Adds default record_stream() to ExtraInputs and StoppingCriteria base classes.
lmdeploy/pytorch/multimodal/data_type.py	Adds MultiModalData.record_stream() for multimodal payload ownership tracking.
lmdeploy/pytorch/model_inputs.py	Adds record_stream() to VisionModelInputs, ModelInputsDelta, and ModelInputs.
lmdeploy/pytorch/engine/model_agent/agent.py	Adds _record_forward_input_stream() and uses it after H2D event wait; centralizes H2D key list.
lmdeploy/pytorch/engine/logits_process.py	Adds SamplingInputs.record_stream() for forward-stream registration.

---

💡 Add Copilot custom instructions for smarter, more guided reviews. Learn how to get started.