mlazos

---
  CUDA Streams + torch.compile: GitHub Issues

  Feature Requests / Support Gaps

  #: #118204 (https://github.com/pytorch/pytorch/issues/118204)
  Title: streams x torch.compile: stream is treated as None sometimes
  Filed: Jan 2024
  By: zou3519
  Status: Closed

mlazos

  ---
  # User code:
  def fn(x):
      s = torch.Stream()
      with s:
          y = x + 1          # runs on side stream
          e.record()          # record_event on side stream
      e.wait()                # wait_event on default stream
      return y * 2            # consumer on default stream

mlazos

"""
User-code repro: control_deps mixed-validity partitioner crash.

Crashes WITHOUT the fix in partitioners.py with:
    AssertionError: Node _unsafe_view was invalid, but is output

Root cause:
  1. s1.wait_stream(default_stream) generates a wait_stream op in the graph.
  2. _collect_wait_stream_forward_deps scans ALL subsequent ops on stream s1 --
     including backward ops -- and collects their inputs defined before the

mlazos

"""
User-code repro: control_deps mixed-validity partitioner crash.

Crashes WITHOUT the fix in partitioners.py with:
    AssertionError: Node _unsafe_view was invalid, but is output

Root cause:
  1. s1.wait_stream(default_stream) generates a wait_stream op in the graph.
  2. _collect_wait_stream_forward_deps scans ALL subsequent ops on stream s1 --
     including backward ops -- and collects their inputs defined before the

mlazos

#!/usr/bin/env python3
"""
Multi-agent code review tool.

Runs 4 review agents in parallel (2 Claude, 2 Codex):
  - Claude high-level design reviewer
  - Codex high-level design reviewer
  - Claude low-level bug finder (real bugs only)
  - Codex low-level bug finder (real bugs only)

mlazos

CUDA Streams in torch.compile
We've added native CUDA stream support to torch.compile, enabling multi-stream GPU programs to be compiled and optimized automatically.
Why it matters: CUDA streams are the primary mechanism for achieving concurrency on GPUs. A stream is a queue of operations that execute in order relative to each other, but can run concurrently with operations on other streams. This enables critical performance patterns like overlapping computation with communication in distributed training, pipelining microbatches across pipeline stages, and offloading activations to CPU during the forward pass and prefetching them during the backward pass. Until now, these patterns required eager execution or careful manual optimization that bypassed the compiler entirely.
Tracing and representing streams in TorchDynamo
TorchDynamo traces torch.cuda.stream() context managers by maintaining a symbolic stream stack — pushing on entry, popping on exit — and annotating every FX graph node with a stream_idx metadat

mlazos

[DEBUG]:Output code: 
# AOT ID: ['0_inference']
from ctypes import c_void_p, c_long, c_int
import torch
import math
import random
import os
import tempfile
from math import inf, nan
from cmath import nanj

mlazos

with torch.Stream(device='cuda') as s_cuda:
    a = torch.randn(10, 5, device='cuda')
    b = torch.randn(5, 10, device='cuda')
    c = torch.mm(a, b)

mlazos

Here's the repro:

  import torch
  torch._dynamo.config.capture_scalar_outputs = True

  def fn(x, val_tensor):
      val = val_tensor.item()  # Creates an unbacked float (fp64)
      scaled = val * 2.0       # fp64 * fp64 = fp64 computation
      return x * scaled        # fp32 * fp64 -> needs downcast

mlazos

import torch
from torch._dynamo.testing import AotEagerAndRecordGraphs
import torch.fx.traceback as fx_traceback

def forward(x):
    with fx_traceback.annotate({"pp_stage": 0}):
        with fx_traceback.annotate({"fdsp_bucket": 0}):
            sin = torch.sin(x)
        sub = sin - 2
    with fx_traceback.annotate({"cuda_stream": 2, "fsdp_bucket": 1}):