Yuhao Ge

I am an LLM infrastructure and performance engineer with rare cross-accelerator expertise (TPU, Trainium, GPU). My work spans the full LLM infrastructure stack, from compiler optimization, kernel implementation, to large-scale training and inference systems.

Currently, I am working at Google, where I focus on improving Gemini training and serving performance on TPU, and enhancing third party TPU experience through TorchTPU.

Previously, I worked at Amazon AWS, Annapurna Labs, where I focused on LLM performance optimization on Trainium chips, primarily working on compiler optimization, kernel language design, and kernel optimization.

I was a Master’s student in Computer Science at UIUC, working under Prof. Charith Mendis. My research interests span Machine Learning, Compilers, and LLM Efficiency.

I earned dual B.S. degrees in Computer Engineering from Zhejiang University and the University of Illinois at Urbana-Champaign through their joint program. I also was a visiting research student at the UCLA VAST Lab, under the supervision of Prof. Jason Cong, where I worked on FPGA accelerator design automation. Additionally, I have interned at TikTok and NFTGo.