Seamless Integration: Custom Kernels in the Torch.compile Stack Wi...
summary
Transformer Engine integration with torch.compile focuses on eliminating graph breaks by registering kernels as first-class ops so Dynamo sees output metadata and aliases without falling back to eager. Practical steps: wrap kernels with custom_op decorator, supply a fake kernel for tracing, declare mutate_args and autograd rules, and register custom Python configs via the experimental opaque objects API. Result: dozens of graph breaks become a single graph, lost fusion and guard overhead drop (example: 17 vs 8 guard checks), enabling CUDA-graph reduced-overhead paths and the fastest end-to-end runs.
AI-written summary · about
chapters (16)
- 0:00Custom kernel integration
- 1:49Why Transformer Engine is fast
- 3:14Torch.compile acceleration
- 5:01Problem with graph breaks
- 6:25Graph breaks hurt performance
- 8:10Guard checks overhead
- 9:37Lost fusion from breaks
- 10:20TE layers and graph breaks
- 11:46Custom op decorator
- 12:48In-place semantics in ops
- 13:51Opaque objects in Torch.compile
- 15:15Opaque objects API
- 16:43Autograd for custom ops
- 17:45Performance results
- 19:34Integration progress
- 20:41Q&A on allocation
related talks
09mTogether Goes Brrr: Threading Research & Production with Torch Compile
23mBlazing Fast GenAI Inference With Torch.compile
36mDL Compiler Panel Discussion
12mLightning Talk: Accelerating Inference on CPU with Torch.Compile
15mPyTorch 2.0: Unlocking the Power of Deep Learning with the Torch Compile API
25mMojo + PyTorch: A Simpler, Faster Path To Custom Kernels
24mInside MoE Optimization: A Profiler-Guided Tour of torch.compile and vLLM
24mAdvancements in High-Performance AI/ML through PyTorch's Python Compiler