ADAPT

Overview

As an undergraduate researcher in ADAPT, I work with a small research team studying how machine learning workloads run on specialized accelerator hardware. My work focuses on manually writing and profiling NKI kernels to understand optimization patterns, correctness issues, and hardware constraints that can later inform automatic compiler passes.

What I did

• Wrote and tested custom GEMM kernels in NKI to study how tiling, loop ordering, memory reuse, and blocking affect performance on Trainium.
• Profiled kernels on real Trainium hardware with NKI profiling tools and Perfetto traces to understand compute activity, memory movement, and bottlenecks.
• Investigated Trainium-specific constraints around PSUM, SBUF, and matrix multiplication behavior that affected kernel correctness and performance.
• Designed accumulation and blocking strategies for cases where larger tensor shapes could not fit cleanly into hardware accumulator limits.
• Documented findings across kernel variants so the results could inform future compiler passes for tiling, loop ordering, and memory reuse.

Why this matters

I'm drawn to this work because low-level systems research can change who gets to build with AI. Better compiler and accelerator tooling can make models cheaper and easier to run, which lowers the barrier for researchers, students, and smaller teams without massive compute budgets.

Stack