Yuhao Ge

I am an LLM serving and performance engineer with cross-accelerator expertise (TPU, Trainium, GPU). My work spans framework, compiler, and kernel optimization for large-scale serving systems.

At Google, I lead development of the PyTorch-native vLLM TPU backend, making TPU a drop-in serving target with GPU-equivalent usability and performance. I work across the full TPU serving stack, including asynchronous scheduling, eager and compiled execution, static-shape bucketing, distributed serving, Pallas kernels, quantization, and TorchTPU. I also bring up and optimize open-source models such as GPT-OSS, Qwen3, and Qwen3.5 for production workloads on TPU.

Previously, I was a founding engineer of the NKI kernel language and compiler at AWS Annapurna Labs. I designed and optimized compiler passes, custom FlashAttention kernels, and autotuning infrastructure for Trainium.

I earned an M.S. in Computer Science and a B.S. in Computer Engineering from UIUC, graduating with Highest Honors and Bronze Tablet recognition. During my master’s research with Prof. Charith Mendis, I developed SPLAT, an optimized GPU code-generation framework for sparse attention accepted at OOPSLA 2025.

Earlier, I worked on ML-driven FPGA accelerator optimization at the UCLA VAST Lab with Prof. Jason Cong, and on real-time avatar animation systems at TikTok.