Add review comment

Author: nvcastet

python/sglang/srt/distributed/device_communicators/pynccl_allocator.py

@@ -1,6 +1,7 @@
 import tempfile
 
 import torch
+from packaging import version
 from torch.cuda.memory import CUDAPluggableAllocator
 
 from sglang.srt.distributed.parallel_state import GroupCoordinator
@@ -105,7 +106,9 @@ def __exit__(self, exc_type, exc_val, exc_tb):
         self._mem_pool_ctx.__exit__(exc_type, exc_val, exc_tb)
         for segment in get_nccl_mem_pool().snapshot():
             if segment["address"] not in _registered_base_addrs:
-                if segment["stream"] == 0:
+                if segment["stream"] == 0 and version.parse(
+                    torch.__version__
+                ) < version.parse("2.8.0"):
                     # PyTorch version < 2.8.0 has a multi-thread MemPool bug
                     # See https://github.com/pytorch/pytorch/issues/152861
                     # Fixed at https://github.com/pytorch/pytorch/commit/f01e628e3b31852983ab30b25bf251f557ba9c0b