work

Amazon, Annapurna Labs
Software Engineer Intern
2024.5 - 2024.8

Worked with the Neuron Compiler team, focusing on optimizing deep learning on AWS AI accelerators (Trainium/Inferentia).

  • Developed infrastructure for automatic kernel generation, compilation, profiling, and visualization, defining a comprehensive sweep space.

  • Collected data for DMA access pattern analysis and introduced a learning-based DMA latency model.

  • Created the first-generation autotuning infrastructure from scratch, enabling compiler and kernel optimizations.

  • Applied autotuning to optimize the Matrix Multiply Fusion Pass, achieving a 14.7% performance improvement for the Llama3.1 model.

  • Developed a kernel language for AI accelerators and supported kernel optimization through autotuning, leading to a 4.9% HFU improvement for kernels like Matrix Multiply.

  • Implemented multi-process compilation and distributed benchmarking, resulting in an 8.62X speedup.

  • Delivered three organization-wide presentations and received a Certificate of Appreciation from the director for developing one of the most impactful internship projects, significantly enhancing peak performance on Trainium and pioneering new avenues for future development.

NFTGo
Machine Learning Engineer Intern
2023.2 - 2023.6

Built an NFT pricing service powered by machine learning

  • Developed a Regression Model for NFT pricing using historical transaction data and NFT features

  • Used MongoDB and Redis to realize efficient data retrieval, and the FastAPI web framework to package the API services

  • Streamlined periodic data processing, model training/updating, and monitoring with Apache Airflow

  • Deployed and managed the system using Docker and Kubernetes and monitored through Grafana

TikTok
Software Engineer Intern
2022.5 - 2022.8

Collaborated in developing TikTok 3D Game Engine, which empowers users to create/use interactive AR/VR stickers

  • Implemented a query-based animation system Motion Matching in C++ for realistic and responsive avatar control

  • Developed an SDK for Skeleton Retargeting in C++ and Lua, supporting animation adaptation across character models

  • Integrated the cross-functional team’s Text-to-Animation algorithm into our game engine using the developed SDK