How does async copy with/without streams work?

dmadeka1: <p>I’m trying to run this bit of code, and expecting the final matrix multiply to work asynchronously while the copies happen. Im surprised that the trace actually shows the operations to be synchronous</p>
<p>Even if I remove the streams, the same trace occurs.</p>
<pre data-code-wrap="python"><code class="lang-python">import torch
import time

stream1 = torch.cuda.Stream()
stream2 = torch.cuda.Stream()
stream3 = torch.cuda.Stream()

N = 1000
mul = 5
num = 1000

qs = []
for i in range(num):
    qs.append(torch.rand(N, N, device='cuda:3'))
q2 = torch.rand(mul * N, mul * N, device='cuda:5')

rs = []
for i in range(num):
    rs.append(torch.rand(N, N, device='cuda:4'))
r2 = torch.rand(mul * N, mul * N, device='cuda:6')

from torch.profiler import profile, record_function, ProfilerActivity

## Warm up

for q in qs:
    q @ q

for r in rs:
    r @ r

q2 = q2 @ q2
r2 = r2 @ r2

for i in ['3', '4', '5', '6']:
    temp = 'cuda:{}'.format(i)
    print(temp)
    torch.cuda.synchronize(torch.device(temp))

with profile(activities=[ProfilerActivity.CPU, ProfilerActivity.CUDA]) as prof2:
    for q in qs[:len(qs) // 2]:
        q.to('cuda:5')
    q2 = q2 @ q2 
    q2 = q2 @ q2

    for r in rs[:len(qs) // 2]:
        r.to('cuda:6', non_blocking=True)
    r2 = r2 @ r2 
    r2 = r2 @ r2

    for i in ['3', '4', '5', '6']:
        temp = 'cuda:{}'.format(i)
        print(temp)
        torch.cuda.synchronize(torch.device(temp))

    with torch.cuda.stream(stream1):
        for q in qs[len(qs) // 2:]:
            q.to('cuda:5')

        q2 = q2 @ q2 
        q2 = q2 @ q2
        
    with torch.cuda.stream(stream2):
        r2 = r2 @ r2 
        r2 = r2 @ r2
            
    with torch.cuda.stream(stream3):
        for r in rs[len(qs) // 2:]:
            r.to('cuda:6', non_blocking=True)

prof2.export_chrome_trace('check12.json')
</code></pre>
<p>The trace seems to shows concurrency for both sets of workloads, which very much surprised me and made me think I dont really understand :(.</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/f/c/fca1cec04dcdaa5729cfb8366d0ce7eef822e0bc.png" data-download-href="https://discuss.pytorch.org/uploads/default/fca1cec04dcdaa5729cfb8366d0ce7eef822e0bc" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fca1cec04dcdaa5729cfb8366d0ce7eef822e0bc_2_690x188.png" alt="image" data-base62-sha1="A2SZC6zDpvsgWIhsPilvpdmYzDu" width="690" height="188" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fca1cec04dcdaa5729cfb8366d0ce7eef822e0bc_2_690x188.png, https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fca1cec04dcdaa5729cfb8366d0ce7eef822e0bc_2_1035x282.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/f/c/fca1cec04dcdaa5729cfb8366d0ce7eef822e0bc_2_1380x376.png 2x" data-dominant-color="CDDCDC"><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">image</span><span class="informations">2408×658 54.8 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>Zooming in to GPU:4 and GPU:6</p>
<p><div class="lightbox-wrapper"><a class="lightbox" href="https://discuss.pytorch.org/uploads/default/original/3X/9/4/94307a746c7f49452a0902760a53e70b84e2ee9f.png" data-download-href="https://discuss.pytorch.org/uploads/default/94307a746c7f49452a0902760a53e70b84e2ee9f" title="image"><img src="https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94307a746c7f49452a0902760a53e70b84e2ee9f_2_690x205.png" alt="image" data-base62-sha1="l8WwRh9pVEL0IF1I3orYu1xpET5" width="690" height="205" srcset="https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94307a746c7f49452a0902760a53e70b84e2ee9f_2_690x205.png, https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94307a746c7f49452a0902760a53e70b84e2ee9f_2_1035x307.png 1.5x, https://discuss.pytorch.org/uploads/default/optimized/3X/9/4/94307a746c7f49452a0902760a53e70b84e2ee9f_2_1380x410.png 2x" data-dominant-color="D3DBDE"><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">image</span><span class="informations">2420×720 110 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>Im surprised that stream 3 waits on stream 2 (even if I flip the order I get the same behavior). Removing the stream contexts didn’t really generate the async behavior I expected either <img src="https://discuss.pytorch.org/images/emoji/apple/frowning.png?v=12" title=":frowning:" class="emoji" alt=":frowning:" loading="lazy" width="20" height="20"></p>
<p>Dhruv</p>