Default PyTorch CUDA stream copy ops surprisingly block other copy ops from other PyTorch non-blocking CUDA stream

liuzijing2014: <p>Hi experts,</p>
<p>We are working on LLM inference, and notice something surprising from our model inference stack. From our observation, the copy operations (e.g. <code>pytorch.tensor.to()</code>/<code>cpu()</code>/<code>tolist()</code>, for both <code>non_blocking=True/False</code>) would block other copy operations running on other CUDA streams whose <code>cudaStreamNonBlocking</code> is set.</p>
<p>Per <a href="https://docs.nvidia.com/cuda/cuda-runtime-api/stream-sync-behavior.html" rel="noopener nofollow ugc">CUDA doc</a>, we are not expecting this because it seems that default CUDA stream only blocks other CUDA streams who are not non-blocking.</p>
<p>Here is a quick reproduce script:</p>
<pre><code class="lang-auto">import ctypes

import threading
from ctypes import c_int, c_uint, c_void_p

import torch

# Load CUDA runtime library
try:
    cuda_runtime = ctypes.CDLL("libcudart.so")  # Linux
except OSError:
    try:
        cuda_runtime = ctypes.CDLL("cudart64_*.dll")  # Windows
    except OSError:
        cuda_runtime = ctypes.CDLL("libcudart.dylib")  # macOS


# Define CUDA constants
cudaStreamNonBlocking = 0x01


def is_stream_non_blocking(pytorch_stream: torch.cuda.Stream) -&gt; bool:
    """Check if a PyTorch CUDA stream is non-blocking."""
    try:
        # Get the raw CUDA stream handle
        stream_ptr = pytorch_stream.cuda_stream

        # Prepare arguments for cudaStreamGetFlags
        flags = c_uint()

        # Call cudaStreamGetFlags
        result = cuda_runtime.cudaStreamGetFlags(
            c_void_p(stream_ptr), ctypes.byref(flags)
        )

        if result != 0:  # cudaSuccess = 0
            print(f"Error getting stream flags: {result}")
            return False

        return bool(flags.value &amp; cudaStreamNonBlocking)

    except Exception as e:
        print(f"Error checking stream flags: {e}")
        return False


def h2d_func():
    stream_h2d = torch.cuda.Stream()
    t = torch.ones([10000], device=torch.device("cpu"), pin_memory=True)
    with torch.cuda.stream(stream_h2d):
        print(f"stream_h2d is non-blocking: {is_stream_non_blocking(stream_h2d)}")
        for i in range(100000000):
            cuda_t = t.to("cuda", non_blocking=True)


def main() -&gt; None:
    h2d_thread = threading.Thread(target=h2d_func, args=())
    h2d_thread.start()
    list_t = torch.empty([1])
    cuda_t = torch.arange(1000000, device=torch.device("cuda"))
    stream_main = torch.cuda.Stream()
    print(
        f"Current stream is non-blocking: {is_stream_non_blocking(torch.cuda.current_stream())}"
    )
    # with torch.cuda.stream(stream_main):
    for i in range(10000000):
        for j in range(100):
            t = cuda_t.repeat_interleave(10)
            t = cuda_t.repeat_interleave(10)
        list_t = t.to('cpu', non_blocking=True)

    h2d_thread.join()
    print("h2d thread is done")


if __name__ == "__main__":
    main()

</code></pre>
<p>Result:</p>
<pre><code class="lang-auto">stream_h2d is non-blocking: True
Current stream is non-blocking: False
</code></pre>
<p>From the GPU trace, we can confirm that there is some blocking issue:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/b/b/bbb191ffbd8b65c4b1c5d1dc07cbdfd1d55de837.png" data-download-href="https://discuss.pytorch.org/uploads/default/bbb191ffbd8b65c4b1c5d1dc07cbdfd1d55de837" title="Screenshot 2025-08-01 at 3.36.25 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbb191ffbd8b65c4b1c5d1dc07cbdfd1d55de837_2_690x323.png" alt="Screenshot 2025-08-01 at 3.36.25 PM" data-base62-sha1="qMpJeBp3oF1yYUd6yPFpOQJu7rN" width="690" height="323" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbb191ffbd8b65c4b1c5d1dc07cbdfd1d55de837_2_690x323.png, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbb191ffbd8b65c4b1c5d1dc07cbdfd1d55de837_2_1035x484.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/b/b/bbb191ffbd8b65c4b1c5d1dc07cbdfd1d55de837_2_1380x646.png 2x" data-dominant-color="9B9AA1"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2025-08-01 at 3.36.25 PM</span><span class="informations">1519×713 102 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>
<p>If we change to let all copy operations run in non-default CUDA streams, we can observe that all copy operations could be run in parallel.</p>
<pre><code class="lang-auto">    with torch.cuda.stream(stream_main):
        for i in range(10000000):
            for j in range(100):
                t = cuda_t.repeat_interleave(10)
                t = cuda_t.repeat_interleave(10)
            list_t = t.to('cpu', non_blocking=True)
</code></pre>
<p>My question is that what is causing the blocking here?</p>

liuzijing2014: <p>Trace after using separate CUDA stream:</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/7/2/727c4c1754d61e8d4de49b57e4817102ee0a6e88.png" data-download-href="https://discuss.pytorch.org/uploads/default/727c4c1754d61e8d4de49b57e4817102ee0a6e88" title="Screenshot 2025-08-01 at 3.36.18 PM"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/727c4c1754d61e8d4de49b57e4817102ee0a6e88_2_690x335.png" alt="Screenshot 2025-08-01 at 3.36.18 PM" data-base62-sha1="gkMJHSgVEMwIlCEpnhRCWxFRCDS" width="690" height="335" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/727c4c1754d61e8d4de49b57e4817102ee0a6e88_2_690x335.png, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/727c4c1754d61e8d4de49b57e4817102ee0a6e88_2_1035x502.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/7/2/727c4c1754d61e8d4de49b57e4817102ee0a6e88_2_1380x670.png 2x" data-dominant-color="ABA4AA"><div class="meta"><svg class="fa d-icon d-icon-far-image svg-icon" aria-hidden="true"><use href="#far-image"></use></svg><span class="filename">Screenshot 2025-08-01 at 3.36.18 PM</span><span class="informations">2500×1214 359 KB</span><svg class="fa d-icon d-icon-discourse-expand svg-icon" aria-hidden="true"><use href="#discourse-expand"></use></svg></div></a></div></p>