Enabling State-of-the-art Asynchronous Execution in Torch.compile With CUDA Streams
summary
Adds native CUDA stream support to torch.compile (Dynamo → AOT Autograd → TorchInductor) so eager-style stream context managers become symbolic stack state, stream/event table entries and IR ops, with backward-pass stream propagation, inserted record/wait synchronization, fusion and buffer-reuse restrictions, and a sync dealloc op for safe cross-stream deallocation. Enables microbatch overlapping and activation offloading with measured results: up to ~30–40% peak memory savings at large model scale while runtime overhead drops as model compute increases.
AI-written summary · about
chapters (15)
- 0:00Introduction
- 1:02CUDA streams and use cases
- 2:30Stream integration architecture
- 3:31Tracking streams in Torch Dynamo
- 4:55Streams as graph inputs
- 6:59AOT Autograd sync
- 9:04Preventing node reordering
- 10:50Functionalization and mutation
- 14:08Stream-aware kernel fusion
- 14:49Cross-stream memory safety
- 17:15Microbatch overlapping
- 18:19Activation offloading
- 19:21Activation offloading results
- 20:25Availability and eager mode
- 21:18Q&A
related talks
15mPyTorch 2.0: Unlocking the Power of Deep Learning with the Torch Compile API
23mBlazing Fast GenAI Inference With Torch.compile
09mTogether Goes Brrr: Threading Research & Production with Torch Compile
12mLightning Talk: Accelerating Inference on CPU with Torch.Compile
36mDL Compiler Panel Discussion
15mLightning Talk: AOTInductor: Ahead-of-Time Compilation for PT2 Exported Models
34mSpeeding up Your ML Workload: Pytorch Compile and Distributed Training
11mLightning Talk: Tensor and 2D Parallelism