work

Improving Gemini training and serving performance on TPU Improving third party TPU experience by TorchTPU

Worked at Annapurna Labs, targeting LLM training/inference acceleration on Trainium chips. Primarily worked in ML Compiler Optimization, Kernel Language Design, Kernel Optimization.

Developed the first-generation ML-based cost model and autotuning infra for Compiler and Kernels. Significantly improved the training and inference performance (14.7% for Llama3.1) on the Trainium chips. Recognized as one of the most impactful internship and earned a Certificate of Appreciation.

Built an NFT pricing service powered by machine learning. Developed regression models using historical transaction data and NFT features. Deployed the system using FastAPI, MongoDB, Redis, Apache Airflow, Docker, and Kubernetes.

Developed TikTok 3D Game Engine for interactive AR/VR stickers. Implemented Motion Matching animation system in C++ for realistic avatar control. Built SDK for Skeleton Retargeting in C++ and Lua, and integrated Text-to-Animation algorithms.