From 7c6838e181585bd89dafb6c98fe8fe81bd6424e1 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 6 Jan 2026 09:45:25 -0800 Subject: [PATCH 1/3] activate NCCL_SYMMETRIC auto-tuning Signed-off-by: Ludwig Schneider --- tensorrt_llm/_torch/custom_ops/torch_custom_ops.py | 5 ++--- 1 file changed, 2 insertions(+), 3 deletions(-) diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 5b683637c6e3..90cf4b6b2b54 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1695,8 +1695,7 @@ def get_valid_tactics( **kwargs, ) -> List[int]: valid_strategies = [ - # TODO: NCCL_SYMMETRIC will cause hang during tuning process - # AllReduceStrategy.NCCL_SYMMETRIC.value, + AllReduceStrategy.NCCL_SYMMETRIC.value, AllReduceStrategy.NCCL.value, ] # Fallback in allreduceOp is set to NCCL_SYMMETRIC as default @@ -1725,7 +1724,7 @@ def forward( input, residual, norm_weight, scale, bias, workspace = inputs if tactic == -1: # TODO: Use NCCL instead of NCCL_SYMMETRIC to avoid hanging during tuning process - tactic = AllReduceStrategy.NCCL.value + tactic = AllReduceStrategy.NCCL_SYMMETRIC.value return torch.ops.trtllm.allreduce( input, From bc377bb069609858aaf9157fa51cbe08972037a6 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Wed, 7 Jan 2026 12:44:01 -0600 Subject: [PATCH 2/3] removing todo Signed-off-by: Ludwig Schneider --- .../custom_ops/tensor_lifetime_workaround.py | 127 ++++++++++++++++++ .../_torch/custom_ops/torch_custom_ops.py | 1 - 2 files changed, 127 insertions(+), 1 deletion(-) create mode 100644 tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py diff --git a/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py b/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py new file mode 100644 index 000000000000..11637b0d972e --- /dev/null +++ b/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py @@ -0,0 +1,127 @@ +""" +Workaround for PyTorch tensor lifetime issue during CUDA graph capture. + +This module provides a mechanism to keep tensors alive during CUDA graph capture +by storing Python references to them. This prevents premature destruction of tensors +returned from custom C++ operators that have custom deleters. + +Issue: During CUDA graph capture, tensors returned from custom C++ operators +may have use_count=1, causing them to be destroyed immediately before PyTorch +binding can increment the reference count. This causes custom deleters to be +called prematurely, releasing buffers while NCCL operations are still using them. + +Workaround: Store references to output tensors during graph capture, ensuring +they stay alive until graph execution completes. +""" + +import threading +from typing import List, Union + +import torch + + +class TensorLifetimeRegistry: + """ + Thread-safe registry to store tensor references during CUDA graph capture. + + This ensures tensors returned from custom operators stay alive during graph + capture and execution, preventing premature deleter calls. + """ + + def __init__(self): + self._lock = threading.Lock() + # Store references per thread to handle multi-threaded scenarios + self._thread_local = threading.local() + + def _get_storage(self) -> List[List[torch.Tensor]]: + """Get thread-local storage for tensor references.""" + if not hasattr(self._thread_local, "tensor_refs"): + self._thread_local.tensor_refs = [] + return self._thread_local.tensor_refs + + def register_tensors(self, tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): + """ + Register tensor(s) to keep them alive during graph capture. + + Args: + tensors: Single tensor, list of tensors, or tuple of tensors to register + """ + with self._lock: + storage = self._get_storage() + + # Convert to list of tensors + if isinstance(tensors, torch.Tensor): + tensor_list = [tensors] + elif isinstance(tensors, (list, tuple)): + tensor_list = [t for t in tensors if isinstance(t, torch.Tensor)] + else: + return # Not a tensor, ignore + + # Only register if we're in graph capture + if self.is_capturing(): + storage.append(tensor_list) + print( + f"[TensorLifetimeRegistry] Registered {len(tensor_list)} tensor(s) " + f"during graph capture (total batches: {len(storage)})" + ) + + def is_capturing(self) -> bool: + """ + Check if we're currently in CUDA graph capture. + + Returns: + True if currently capturing a CUDA graph, False otherwise + """ + try: + # Check if any stream is currently capturing + # torch.cuda.is_current_stream_capturing() checks the current stream + return torch.cuda.is_current_stream_capturing() + except (AttributeError, RuntimeError): + # Fallback: if the function doesn't exist or there's an error, assume not capturing + return False + + def clear(self): + """Clear all registered tensor references (call after graph execution completes).""" + with self._lock: + if hasattr(self._thread_local, "tensor_refs"): + count = sum(len(batch) for batch in self._thread_local.tensor_refs) + self._thread_local.tensor_refs.clear() + print(f"[TensorLifetimeRegistry] Cleared {count} tensor reference(s)") + + def get_registered_count(self) -> int: + """Get the number of registered tensor batches.""" + with self._lock: + if hasattr(self._thread_local, "tensor_refs"): + return len(self._thread_local.tensor_refs) + return 0 + + +# Global singleton instance +_tensor_registry = TensorLifetimeRegistry() + + +def register_tensor_references(tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): + """ + Register tensor(s) to keep them alive during CUDA graph capture. + + This is a convenience function that uses the global registry. + + Args: + tensors: Single tensor, list of tensors, or tuple of tensors to register + """ + _tensor_registry.register_tensors(tensors) + + +def clear_tensor_references(): + """Clear all registered tensor references.""" + _tensor_registry.clear() + + +def is_graph_capturing() -> bool: + """Check if we're currently in CUDA graph capture.""" + return _tensor_registry.is_capturing() + + +def get_registered_count() -> int: + """Get the number of registered tensor batches.""" + return _tensor_registry.get_registered_count() diff --git a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py index 90cf4b6b2b54..0fbb5ceacf70 100644 --- a/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py +++ b/tensorrt_llm/_torch/custom_ops/torch_custom_ops.py @@ -1723,7 +1723,6 @@ def forward( ) -> torch.Tensor: input, residual, norm_weight, scale, bias, workspace = inputs if tactic == -1: - # TODO: Use NCCL instead of NCCL_SYMMETRIC to avoid hanging during tuning process tactic = AllReduceStrategy.NCCL_SYMMETRIC.value return torch.ops.trtllm.allreduce( From 43c3a025d1d146dc6dae02f3e113565d9d186b10 Mon Sep 17 00:00:00 2001 From: Ludwig Schneider Date: Tue, 13 Jan 2026 13:27:40 -0600 Subject: [PATCH 3/3] Delete tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py It was added by mistake. Signed-off-by: Ludwig Schneider --- .../custom_ops/tensor_lifetime_workaround.py | 127 ------------------ 1 file changed, 127 deletions(-) delete mode 100644 tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py diff --git a/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py b/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py deleted file mode 100644 index 11637b0d972e..000000000000 --- a/tensorrt_llm/_torch/custom_ops/tensor_lifetime_workaround.py +++ /dev/null @@ -1,127 +0,0 @@ -""" -Workaround for PyTorch tensor lifetime issue during CUDA graph capture. - -This module provides a mechanism to keep tensors alive during CUDA graph capture -by storing Python references to them. This prevents premature destruction of tensors -returned from custom C++ operators that have custom deleters. - -Issue: During CUDA graph capture, tensors returned from custom C++ operators -may have use_count=1, causing them to be destroyed immediately before PyTorch -binding can increment the reference count. This causes custom deleters to be -called prematurely, releasing buffers while NCCL operations are still using them. - -Workaround: Store references to output tensors during graph capture, ensuring -they stay alive until graph execution completes. -""" - -import threading -from typing import List, Union - -import torch - - -class TensorLifetimeRegistry: - """ - Thread-safe registry to store tensor references during CUDA graph capture. - - This ensures tensors returned from custom operators stay alive during graph - capture and execution, preventing premature deleter calls. - """ - - def __init__(self): - self._lock = threading.Lock() - # Store references per thread to handle multi-threaded scenarios - self._thread_local = threading.local() - - def _get_storage(self) -> List[List[torch.Tensor]]: - """Get thread-local storage for tensor references.""" - if not hasattr(self._thread_local, "tensor_refs"): - self._thread_local.tensor_refs = [] - return self._thread_local.tensor_refs - - def register_tensors(self, tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): - """ - Register tensor(s) to keep them alive during graph capture. - - Args: - tensors: Single tensor, list of tensors, or tuple of tensors to register - """ - with self._lock: - storage = self._get_storage() - - # Convert to list of tensors - if isinstance(tensors, torch.Tensor): - tensor_list = [tensors] - elif isinstance(tensors, (list, tuple)): - tensor_list = [t for t in tensors if isinstance(t, torch.Tensor)] - else: - return # Not a tensor, ignore - - # Only register if we're in graph capture - if self.is_capturing(): - storage.append(tensor_list) - print( - f"[TensorLifetimeRegistry] Registered {len(tensor_list)} tensor(s) " - f"during graph capture (total batches: {len(storage)})" - ) - - def is_capturing(self) -> bool: - """ - Check if we're currently in CUDA graph capture. - - Returns: - True if currently capturing a CUDA graph, False otherwise - """ - try: - # Check if any stream is currently capturing - # torch.cuda.is_current_stream_capturing() checks the current stream - return torch.cuda.is_current_stream_capturing() - except (AttributeError, RuntimeError): - # Fallback: if the function doesn't exist or there's an error, assume not capturing - return False - - def clear(self): - """Clear all registered tensor references (call after graph execution completes).""" - with self._lock: - if hasattr(self._thread_local, "tensor_refs"): - count = sum(len(batch) for batch in self._thread_local.tensor_refs) - self._thread_local.tensor_refs.clear() - print(f"[TensorLifetimeRegistry] Cleared {count} tensor reference(s)") - - def get_registered_count(self) -> int: - """Get the number of registered tensor batches.""" - with self._lock: - if hasattr(self._thread_local, "tensor_refs"): - return len(self._thread_local.tensor_refs) - return 0 - - -# Global singleton instance -_tensor_registry = TensorLifetimeRegistry() - - -def register_tensor_references(tensors: Union[torch.Tensor, List[torch.Tensor], tuple]): - """ - Register tensor(s) to keep them alive during CUDA graph capture. - - This is a convenience function that uses the global registry. - - Args: - tensors: Single tensor, list of tensors, or tuple of tensors to register - """ - _tensor_registry.register_tensors(tensors) - - -def clear_tensor_references(): - """Clear all registered tensor references.""" - _tensor_registry.clear() - - -def is_graph_capturing() -> bool: - """Check if we're currently in CUDA graph capture.""" - return _tensor_registry.is_capturing() - - -def get_registered_count() -> int: - """Get the number of registered tensor batches.""" - return _tensor_registry.get_registered_count()