Revert "已修改 vllm/utils/__init__.py"

This reverts commit 20b6cf64.

vllm/utils/__init__.py

@@ -1173,11 +1173,12 @@ def find_nccl_library() -> str:
 
 prev_set_stream = torch.cuda.set_stream
 
-_current_stream_tls = threading.local()
+_current_stream = None
 
 
 def _patched_set_stream(stream: torch.cuda.Stream) -> None:
-    _current_stream_tls.value = stream
+    global _current_stream
+    _current_stream = stream
     prev_set_stream(stream)
 
 
@@ -1196,8 +1197,8 @@ def current_stream() -> torch.cuda.Stream:
     from C/C++ code.
     """
     from vllm.platforms import current_platform
-    if not hasattr(_current_stream_tls,
-                   "value") or _current_stream_tls.value is None:
+    global _current_stream
+    if _current_stream is None:
         # when this function is called before any stream is set,
         # we return the default stream.
         # On ROCm using the default 0 stream in combination with RCCL
@@ -1207,9 +1208,8 @@ def current_stream() -> torch.cuda.Stream:
         # fix computational precision issue in eager mode 
         # _current_stream = torch.cuda.Stream() if current_platform.is_rocm(
         # ) else torch.cuda.current_stream()
-        _current_stream_tls.value = torch.cuda.Stream(
-        ) if current_platform.is_rocm() else torch.cuda.current_stream()
-    return _current_stream_tls.value
+        _current_stream = torch.cuda.current_stream()
+    return _current_stream
 
 
 def enable_trace_function_call_for_thread(vllm_config: VllmConfig) -> None: