Staff Software Engineer, Machine Learning — Google Cloud
I am a Staff Machine Learning Engineer at Google Cloud with 8 years of industry experience in recommendation systems and large-scale retrieval and ranking, currently focused on open-source LLM inference on TPUs.
Outside of work, I am an avid smart home and IoT enthusiast who enjoys experimenting with all kinds of smart lighting setups at home. When I am unplugged, I love long-distance running and exploring the underwater world as an Advanced Open Water certified scuba diver.
Google Cloud · YouTube · Uber · Amazon
June 2026 — Present
Mountain View, CA
Core ML — LLM Inference on TPUs
Core contributor to the TPU-native KV-cache transfer and offloading library (tpu-sync) and the PyTorch-native TPU backend for SGLang.
Aug 2018 — May 2026
San Bruno, CA
Growth — Recommendations Quality
Tech Lead for YouTube Recommendations ML retrieval, architecting foundational long form Videos and Shorts recommendation models.
May 2017 — Aug 2017
San Francisco, CA
Piper — Workflow Management System
Improved the logging infrastructure of Piper (Uber's Apache Airflow fork) by migrating from MySQL to Cassandra and integrating Elasticsearch, eliminating database write contention and enabling fast, full-text log search at scale.
May 2015 — Jul 2015
Bengaluru, India
Value Added Services (VAS) — Amazon Home Services
Enhanced the seller onboarding experience on the Value Added Services platform by delivering dynamic image transformation tools and custom map-based boundary management for localized service fulfillment.
T. D. Nguyen , D. Jain , M. Lawrence , S. Rendle , K. Su , A. Subbiah
Novel parameter-efficient embedding representations enabling massive item vocabularies within strict accelerator memory budgets.
V. Gatto , J. Li , C. Chen , H. Fu , B. Froehle , S. Hinsu , D. Jain , et al.
System architecture for disaggregating and routing partitioned deep neural network submodels across heterogeneous CPU/GPU/TPU pools.
M. Chen , C. Xu , V. Gatto , D. Jain , A. Kumar , E. Chi
Production-scale reinforcement learning framework introducing a temporal-difference critic network to mitigate off-policy bias and optimize long-term user satisfaction.
D. Jain , M. Kustikova , M. Darbari , R. Gupta , S. Mayhew
Demonstrated strong named entity recognition baselines on noisy multilingual code-switched social media streams.
google/tpu-sync — TPU-Native KV Cache Transfer & Offloading
Core contributor to Google’s open-source C++ and Python library for high-performance Key-Value (KV) cache orchestration, multi-tiered memory management (HBM/DRAM), and zero-copy cross-node DMA data transfers on Google Cloud TPUs.
PyTorch-Native TPU Backend for SGLang
Core contributor to the PyTorch-native TPU backend and plugin for SGLang, enabling high-concurrency LLM and VLM serving on Google TPUs with unified KV-cache pooling, prefix caching, and disaggregated prefill/decode execution.
Fault-Tolerant Distributed Cloud Platform
Developed a fault-tolerant distributed cloud platform in C++ featuring a front-end load balancer, a multithreaded TCP-based HTTP server with CGI script execution, SMTP (with relay) and POP3 email servers, and a distributed storage service supporting file and folder creation, movement, renaming, and deletion.
Replicated Multi-Node Multicast Chat Server
Built a replicated distributed chat server in C++ supporting concurrent multi-room communication across nodes using UDP multicast, implementing unordered, FIFO, and total ordering message delivery protocols.
Deep Learning-Based Text Simplification
Built a neural text simplification pipeline using attentional bidirectional RNN sequence-to-sequence (Seq2Seq) models on the Newsela dataset, combining named entity anonymization, diverse beam-search clustering, and LSTM language model perplexity reranking evaluated with SARI and BLEU.
Personalized Adaptive E-Learning Recommendation Engine
Developed an adaptive e-learning platform and recommendation engine using Python, Django, and Celery that dynamically personalizes educational content and learning pathways based on student proficiency and interaction telemetry.
Healthcare & Clinic Workflow Management Application
Engineered a clinic management application in Java and JavaFX backed by a relational database to manage patient records, doctor scheduling, medical consultations, and automated PDF prescription generation.
Core technical areas across production ML serving and large-scale retrieval.
Accelerator-native training and inference runtimes.
High-availability orchestration and disaggregated execution.
Systems and machine learning languages.
Master of Science in Engineering (M.S.E.)
Computer and Information Science
Bachelor of Technology (B.Tech.)
Information and Communication Technology