[Bugfix] Carry over queued work when materializing the dedicated stream (#56382)
Signed-off-by: zhaoguochun1995 <zhaoguochun1995@163.com>
zhaoguochun1995 committed
Sep 15, 2026 at 07:52 UTC
bbbd0a02c9d1bec965ee818af29a1319baefa69c
1 file changed
+4
-1
vllm/utils/torch_utils.py
+4
-1
@@ -816,8 +816,11 @@ def current_stream() -> torch.cuda.Stream:
816
# https://github.com/pytorch/pytorch/blob/42ad9edfb754743fdae3276ade43de000beb4f60/aten/src/ATen/cuda/CUDAGraph.cpp#L77
817
# for more details. Therefore, we create a dedicated stream per process.
818
if current_platform.is_rocm() or current_platform.is_cuda():
819
+ # Ensure new stream is ordered w.r.t. replaced stream's work.
820
+ new_stream = torch.cuda.Stream()
821
+ new_stream.wait_stream(torch.cuda.current_stream())
822
# torch.cuda.set_stream here is the alias of _pathed_set_stream
820
- torch.cuda.set_stream(torch.cuda.Stream())
823
+ torch.cuda.set_stream(new_stream)
824
elif current_platform.is_cpu():
825
_current_stream_tls.value = _StreamPlaceholder()
826
else: