I am a second-year Master's student in Electrical and Computer Engineering at the University of Texas at Austin, focusing on Computing Systems for Large Scale AI. I am currently working as a research assistant with Dr. Poulami Das, where I am working on improving the LLM inference efficiency driven by fundamental Computer Architecture principles.
I recently completed an internship at Together.ai, working as a researcher on the training team. My work focused on improving LLM training efficiency for long-context scenarios. My proposed method of splitting attention head-wise and pipelining them enabled better memory efficiency for long-context training, allowing to train a Llama-8B model on a single 8xH100 node with upto 5M tokens in context (beating the prior SOTA of 4M tokens).
Prior to this, I worked as a GPU Architect at NVIDIA, Bengaluru, working on GPU Architecture optimization, where I built ML-powered predictive models to project performance for NVIDIA's future products. I also helped build analytical models to debug inefficiencies in the current products, particularly for the gaming market (e.g NVIDIA DLSS). Additionally, I have worked as a part-time researcher at H2Lab, University of Washington with Prof. Prithviraj Ammanabrolu, where I worked on deriving rewards from freeform linguistic feedback to train Large Language Models using RL.
My research interests lie in designing efficient systems for running massive AI models. To this end, I am interested to explore different layers of the system: hardware architecture, compilation techniques, kernel-level software optimizations and library-level (e.g PyTorch) optimizations. I have been fortunate to have an exposure to each of these aspects across different projects during internships as well as coursework. On a tangential note, I am also interested to pursue Reinforcement Learning based finetuning, and interpretability research.
I graduated from IIT Kharagpur in 2021 with a major in Electronics and Elec. Comm. and a minor in Computer Science. My bachelor thesis was on Training Generative Dialog Models using Reinforcement Learning with Learned Reward Systems, supervised by Dr. Pawan Goyal.
My hobbies include trekking, music, cooking (traditional Indian cuisines). I also love reading blogs on abstract ideas, particularly related to tech-stuff and human evolution.