Coding and Reasoning LLMs

At Google I lead AI-4-AI work on post-training, reinforcement learning and agentic workflows. The goal is long-horizon coding agents that automatically discover and maximize the performance of state-of-the-art AI algorithms on accelerators, and AI-assisted tools that make developers more productive.

Research Projects

NeurIPS 2026 Spotlight

MaxKernel

Long-horizon coding agents for TPU kernel optimization, part of our multi-year effort on AI that optimizes the software stack for AI. MaxKernel will be presented as a spotlight at NeurIPS 2026 (MLSys track).

arXiv:2609.04523

ICML DL4Code

JAXBench: Benchmarking Autonomous TPU Kernel Optimization

A TPU performance benchmark that measures how well AI agents can write and tune high-performance JAX/Pallas kernels. JAXBench is now published and was presented at the ICML Workshop (DL4Code).

arXiv:2607.20466

MLE Agent

An agentic framework for migrating deep learning workloads using in-context learning.

arXiv:2606.15994

Gemini Code Assist

Code completion, generation and chat built on the Gemini family of models and integrated into popular IDEs, with enterprise customization through retrieval-augmented generation (RAG).

Coverage: TechCrunch, 2024

Forecasting machine degradation of GPU clusters

Predicting hardware degradation in large GPU clusters before it disrupts training jobs. Machine Learning for Systems, 2025.

OpenReview

Talks

AI System Devlabs presentationSupercharging developer productivity using generative AI (Google Cloud Next)Full codebase awareness (Google Cloud Next)

Selected Publications

  1. Tschand, A., Hong, C., Walker, J., Cai, N., Wang, S., Subramanian, S., Dev, S., Reddi, V.J., Yazdanbakhsh, A. and Sankaran, S. JAXBench: Benchmarking Autonomous TPU Kernel Optimization. ICML Workshop on Deep Learning for Code (DL4Code), 2026. [arXiv]
  2. Liang, Q., Ingram, S., Vanica, G., Gavrilescu, A., Harrat, N., Sipra, H. and Sankaran, S. Agentic Framework for Deep Learning Workload Migration via In-Context Learning. arXiv:2606.15994, 2026. [arXiv]
  3. Cai, S., Nie, S., Chen, Z., Gulalkari, N., Vanica, G., Jain, C. and Sankaran, S. Forecasting machine degradation of GPU clusters. Machine Learning for Systems, 2025. [OpenReview]