Abstract:

This transcript captures a live Q&A session featuring Richard Zoo from the PyTorch compiler team, discussing torch.compile. The session addresses common performance pitfalls such as graph breaks, recompilations, and the handling of non-native custom kernels, offering advice on writing compile-friendly code. Detailed explanations cover integrating third-party kernels, including C++, Python-based operators, and user-defined Trident kernels, highlighting how torch.compile treats them (often as black boxes, with special handling for Trident). Richard outlines the internal architecture (Dynamo for graph capture, AOTDispatcher for normalization/functionalization/autograd, and Inductor for backend optimization/code generation). Key optimizations like operator fusion, dead code elimination, and partitioning are discussed, alongside methods for inspecting the compilation process using logs and the TL parse tool. The future roadmap, focusing on performance, usability (error messages, escape hatches), reducing compile times (including AOT compilation), and expanding library support, is presented. Finally, guidance is given on community contributions to PyTorch and torch.compile.

PyTorch Compile Q&A: Performance, Custom Kernels, Internals, and Future Directions

• 0:00:30 Introduction: Live Q&A session on torch.compile hosted by Mark, featuring Richard Zoo from the PyTorch compiler team. Format involves answering pre-submitted and live chat questions.
• 0:01:43 Common Performance Pitfalls & Optimizable Code: Discussion starts on performance issues with torch.compile and how to write code it can optimize effectively.
• 0:02:09 Pitfall - Graph Breaks: torch.compile breaks the function into subgraphs when it encounters unsupported operations (like print). This incurs overhead and prevents optimizations across breaks. Recommendation: Eliminate breaks or use fullgraph=True to identify them.
• 0:04:58 Pitfall - Custom Operators: If code consists solely of custom C++/CUDA kernels wrapped as operators, torch.compile treats them as black boxes and cannot optimize inside them. Suggestion: Rewrite kernels using native PyTorch ops where possible for better fusion opportunities.
• 0:06:38 Backend Tuning Options: Use mode='reduce_overhead' to enable CUDAGraphs (reduces kernel launch overhead on NVIDIA GPUs). Use autotuning (inductor.config.autotune_pointwise=True, etc.) to let the compiler benchmark and choose the fastest kernel implementation, at the cost of increased compilation time.
• 0:07:52 Pitfall - Recompilations: torch.compile recompiles when inputs change characteristics (e.g., tensor shapes, strides, Python scalar values used in the graph). Excessive recompilation hurts performance. Monitor recompiles and consider using dynamic shapes features if needed.
• 0:09:13 Integrating Custom Kernels: How to use third-party (C++, Python) and user-defined Trident kernels with torch.compile.
• 0:10:05 Kernels vs. Operators: A kernel is the raw computation (e.g., CUDA code). An operator wraps the kernel with PyTorch-specific glue code (metadata, autograd support, dispatch info). Writing an operator is recommended for proper integration.
• 0:11:50 Custom Operator APIs: Use @torch.library.custom_op (Python) or C++ registration APIs. Crucially, provide a meta kernel (or shape inference function) to tell torch.compile how output shapes are derived from input shapes, needed for tracing.
• 0:14:37 Compile's Handling of Custom Ops: Treats them as opaque black boxes. It will execute them but won't optimize their internal logic. This can be used as an "escape hatch" for code that torch.compile cannot handle.
• 0:15:54 Compile's Handling of Trident Kernels: Generally "just work". torch.compile recognizes Trident kernels, wraps them in a black box node, and includes their source code in its generated output. It aims to support most Trident features but sometimes needs updates for new ones (e.g., host TMA, specific autotuning).
• 0:18:41 Trident vs. Custom Op Handling: torch.compile does analyze Trident kernels slightly (unlike regular custom ops) to determine which inputs are mutated (for functionalization) and integrates with Trident's autotuning.
• 0:22:00 Future Roadmap for torch.compile:
• 0:22:52 Continued Performance Improvements: Focus on matmuls (e.g., using Cutlass, targeting new hardware like Blackwell), AMD GPU support.
• 0:23:49 Usability Improvements: Better error messages, better escape hatches for users when compilation fails, making it easier to understand and debug.
• 0:24:46 Compile Time Reduction: Addressing Python overhead in the compiler, exploring ahead-of-time (AOT) compilation workflows to cache compiled artifacts.
• 0:32:33 Expanding Library/Use Case Support: Improving integration with libraries like diffusers, LLM inference/serving frameworks (vLLM, SGLang).
• 0:33:12 Optimizations Performed by torch.compile:
• 0:33:37 Basic Optimizations: Dead code elimination, common subexpression elimination.
• 0:34:52 Min-Cut Partitioning: (For training) Automatically figures out which activations to recompute in the backward pass to save memory, similar to manual activation checkpointing.
• 0:36:19 Pattern Matching: Identifies specific sequences of operations (like manually written attention) and replaces them with optimized implementations (like FlashAttention/SDPA).
• 0:36:42 Kernel Fusion: Combines multiple operations into single kernels. Pointwise ops (like sin, cos, add) and reductions are commonly fused. Use TORCH_LOGS=output_code to inspect.
• 0:40:14 Epilogue Fusion: Fuses elementwise operations that occur after a matmul directly into the matmul kernel.
• 0:41:59 Prologue Fusion: Fuses elementwise operations that occur before a matmul into the matmul kernel (not fully implemented but planned).
• 0:46:18 Matmul Selection/Autotuning: Automatically benchmarks and selects the best matmul implementation (e.g., Triton vs cuBLAS, fused vs unfused) for the specific inputs and hardware.
• 0:48:44 Inspecting Compilation: Use the TL parse tool on logs generated with TORCH_LOGS=+dynamo,+aot,+inductor to view the graph at different stages (Dynamo output, AOT output, Inductor output) and the final generated code in a human-readable format.
• 0:50:56 torch.compile Design Overview:
• 0:57:08 Dynamo (Frontend): Interprets Python bytecode to capture a graph of operations (FX Graph). Handles graph breaks.
• 1:00:21 AOTDispatcher (Middle-end): Normalizes the graph (using ATen operators), performs functionalization (removes mutations), and generates the backward graph using autograd.
• 1:01:17 Inductor (Backend): Takes the functionalized forward/backward graphs, performs further optimizations (fusion, scheduling), and generates low-level code (Triton/C++ for CUDA, C++ for CPU).
• 1:02:34 Training Path: AOTDispatcher generates both forward and backward graphs, Inductor compiles both (backward often lazily).
• 1:06:37 Using Compiled Models in C++: The standard torch.export + AOT Inductor workflow primarily targets inference and produces an artifact loadable in C++ but typically without gradient support.
• 1:07:53 Hacky C++ Training Workflow: Use functorch.make_fx to trace a combined forward-and-backward function into a single graph, then export/compile that graph using AOT Inductor. This artifact can be run in C++, effectively executing both forward and backward.
• 1:11:32 Handling Floating Point Mismatches: Report discrepancies between eager and compiled mode. Compile sometimes uses different (potentially more accurate) numerics. torch._dynamo.config.fast_fp_precision = "higher" might help align behavior. Minimize the compiled region to isolate the source.
• 1:12:36 Contributing to torch.compile: Find "good first issue" or "actionable" issues on GitHub. Read the contribution guide and core onboarding wiki. Engage in discussions on issues/PRs. Help is welcome.
• 1:19:55 Usability Escape Hatches: Plans include allowing specific graph breaks even with fullgraph=True, providing non_strict_trace (a JAX-like tracing mode as an alternative to Dynamo), and mechanisms to override faulty caching behavior.