Skip to content
GitLab
Menu
Projects
Groups
Snippets
Loading...
Help
Help
Support
Community forum
Keyboard shortcuts
?
Submit feedback
Contribute to GitLab
Sign in
Toggle navigation
Menu
Open sidebar
OpenDAS
vllm_cscc
Commits
41088695
Unverified
Commit
41088695
authored
Feb 13, 2025
by
Gregory Shtrasberg
Committed by
GitHub
Feb 14, 2025
Browse files
[ROCm] Avoid using the default stream on ROCm (#13238)
Signed-off-by:
Gregory Shtrasberg
<
Gregory.Shtrasberg@amd.com
>
parent
e38be640
Changes
1
Hide whitespace changes
Inline
Side-by-side
Showing
1 changed file
with
6 additions
and
1 deletion
+6
-1
vllm/utils.py
vllm/utils.py
+6
-1
No files found.
vllm/utils.py
View file @
41088695
...
...
@@ -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
...
...
Write
Preview
Markdown
is supported
0%
Try again
or
attach a new file
.
Attach a file
Cancel
You are about to add
0
people
to the discussion. Proceed with caution.
Finish editing this message first!
Cancel
Please
register
or
sign in
to comment