[Feature]: Check if we can free tensors before capture finalize in PWCG

🚀 The feature, motivation and pitch

Background
In DSR1 we saw a bug due to weak ref storage of the dynamic output tensors. it was fixed in this PR https://github.com/NVIDIA/TensorRT-LLM/pull/13497 by storing strong ref pointers till the end of the capture, then setting them as weak refs.

The Problem
before my fix, during capture, each static segment allocated the next dyanmic segment's outputs and stored it with a weak ref:

            for submod_id, buf in dynamic_out_bufs.items():
                entry.dynamic_out_bufs[submod_id] = make_weak_ref(buf)

That reduced the ref count to zero. Because they're allocated inside a graph-pool capture context, the pool tracks their addresses as "in use" only while Python holds a strong ref. The next static segment read the buffer as input, but since this memory was freed it got allocated inside the segment which caused the corruption (PT is not aware that the input buffer is allocated because it was allocated from the graph pool and used weak ref) and uses it for intermediates. My fix removed the make_weak_ref and set it only after all segments have been captured.

The task: check if we can optimize that by analyzing which buffers can be freed earlier - e.g. after they were already consumed by all consumers

Alternatives

No response

Additional context

No response

Before submitting a new issue...

• Make sure you already searched for relevant issues, and checked the documentation and examples for answers to frequently asked questions.