From 410886950ac5ce773a57ae115a31cdecfcdd4af1 Mon Sep 17 00:00:00 2001 From: Gregory Shtrasberg <156009573+gshtras@users.noreply.github.com> Date: Thu, 13 Feb 2025 20:29:26 -0500 Subject: [PATCH] [ROCm] Avoid using the default stream on ROCm (#13238) Signed-off-by: Gregory Shtrasberg --- vllm/utils.py | 7 ++++++- 1 file changed, 6 insertions(+), 1 deletion(-) diff --git a/vllm/utils.py b/vllm/utils.py index 1d7fbd4a7879..b1bac649c972 100644 --- a/vllm/utils.py +++ b/vllm/utils.py @@ -942,11 +942,16 @@ def current_stream() -> torch.cuda.Stream: the underlying hypothesis is that we do not call `torch._C._cuda_setStream` from C/C++ code. """ + from vllm.platforms import current_platform global _current_stream if _current_stream is None: # when this function is called before any stream is set, # we return the default stream. - _current_stream = torch.cuda.current_stream() + # On ROCm using the default 0 stream in combination with RCCL + # is hurting performance. Therefore creating a dedicated stream + # per process + _current_stream = torch.cuda.Stream() if current_platform.is_rocm( + ) else torch.cuda.current_stream() return _current_stream