gistfile1.txt

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 metadata field that records which stream it executes on. Stream and event operations are lowered to torch.ops.streams.* functions (record_event, wait_event, wait_stream) that take index-based arguments rather than live Python objects.
Since stream and event objects are opaque Python objects that can't be embedded as graph constants, Dynamo generates prefix bytecode that runs before the compiled graph. This prefix populates an index-to-stream table with the live objects, and inside the graph, get_user_obj_by_index reads from this table at runtime to resolve them.
Preserving correctness in AOT Autograd
AOT Autograd preserves stream semantics through functionalization by wrapping event nodes in control_deps higher-order operators. These HOPs take all tensors crossing a stream boundary as inputs and return them as identity passthroughs, creating an explicit data dependency chain in the FX graph that no optimization pass can reorder across.
The backward pass is automatically synchronized: each backward node inherits the stream of its associated forward node, and gradient accumulation nodes are assigned the stream of their first consumer. When gradients from multiple streams feed into the same accumulation, wait_event is inserted before accumulation to ensure all inputs are ready, and record_event plus wait_event broadcast the accumulated result to other streams.
We also detect a subtle interaction with functionalization and input mutations. Functionalization converts in-place ops like x.mul_(2) into functional equivalents, deferring the actual write-back (x.copy_) to a graph epilogue. If an event_record sits between the mutation and the epilogue, the event would fire before the mutation is visible — so Dynamo detects this case at compile time and raises an error, while allowing non-synchronized input mutations where the epilogue ordering doesn't affect correctness.
Stream-aware codegen in TorchInductor
TorchInductor respects stream assignments throughout codegen and scheduling. It emits the necessary record_event and wait_event calls at stream boundaries and enforces fusion constraints: ops on the same stream fuse normally, but ops on different streams are never fused, preserving the concurrency the user intended. Event nodes also act as fusion barriers, preventing ops from being fused across record/wait boundaries.
Cross-stream tensor lifetimes are protected automatically. Dynamo inserts record_event after a tensor's last usage on a side stream and wait_event on the origin stream, which delays allocator reuse of that memory without host-side blocking. The user's synchronize() and del calls are respected — the compiler simply replaces the blocking synchronization with a non-blocking event wait.