4 papers
TorchTPU - A Hybrid Eager/Compiled Architecture for Native PyTorch Execution on TPU
Claudio Basile, Michael Voznesensky, Gunhyun Park +51
TorchTPU is an open-source collaboration between Google and Meta that introduces a native and principled PyTorch backend for Google TPUs. It enables seamless out-of-the-box model e…
AoiZora: Compiler-Agnostic, Placement-Aware Auto-Parallelization for Video Diffusion Transformers
Kaijian Wang, Yuanyuan Xu, Fanjiang Ye +5
Video diffusion transformers (DiTs) process tens of thousands of spatio-temporal tokens per denoising step, so low-latency inference and training commonly spread them across severa…
Ragged Paged Attention: A High-Performance and Flexible LLM Inference Kernel for TPU
Jevin Jiang, Ying Chen, Blake A. Hechtman +2
Large Language Model (LLM) deployment is increasingly shifting to cost-efficient accelerators like Google's Tensor Processing Units (TPUs), prioritizing both performance and total…
Recycling Scraps: Improving Private Learning by Leveraging Intermediate Checkpoints
Virat Shejwalkar, Arun Ganesh, Rajiv Mathews +5
In this work, we focus on improving the accuracy-variance trade-off for state-of-the-art differentially private machine learning (DP ML) methods. First, we design a general framewo…