Subhankar Shah
Subhankar is a Software Engineer on the ML Compilers team at Google, where he works on performance optimizations for large-scale foundation models in the XLA:TPU compiler and JAX ecosystem. His research and engineering interests include accelerator memory management, feedback-driven compiler auto-tuning, and portable ML intermediate representations (StableHLO).
Authored Publications
Sort By
CATWILD: Compiler Autotuning For TPU workloads in the Wild
Kirill Gugaev
Mike Burrows
Mangpo Phothilimthana
Naveen Kumar
Idan Nachman
Nacho Cano
David Liu
Zongwei Zhou
Tengyu Sun
Berkin Ilbeyi
Jie Sun
Vineetha Govindaraj
Ziqiang Feng
Hyojun Kim
Matheus Camargo
Chao Wang
Alex Wertheim
Domagoj Babic
Aleksey Orekhov
Emma Sevastian
Wenqi Cao
Shobha Vasudevan
Karthik Srinivasa Murthy
Jose Baiocchi Paredes
Emma Wang
Ryan Lefever
Arissa Wongpanich
Νikhil Sarda
Mehrdad Khani
Matt Bierbaum
Sami Abu-El-Haija
Samrat Ghosh
2026
Preview abstract
Compilers play a fundamental role at achieving peak performance for machine learning (ML) workloads. However, given the diverse nature of workloads and accelerators, compilers' heuristics and analytical cost models often result in sub-optimal performance, and thus waste precious datacenter resources. Furthermore, the multitude of tunable parameters and their complex interplay often make it impossible for human experts to manually find optimal configurations. In this paper, we present
CATWILD, a system that automatically optimizes ML jobs in Google's TPU fleet using compiler autotuning techniques. We describe CATWILD’s design and implementation, and evaluate its performance using a handful of representative metrics. We further report experiences and lessons learned from its five-year development and operation. To the best of our knowledge, CATWILD represents the first ML compiler autotuning solution deployed in datacenters at scale. Its successful rollout yielded substantial benefits, optimizing over 70\% of daily TPU training jobs and achieving significant chip savings
View details
