Hey! I'm Alex Wa, a 2nd year Math and CS double major at Yale and a YES Scholar. My research interests currently span RL and NLP, and I'm also interested in ML systems and model architecture. Currently, I'm researching RL4LLMs and rubrics as rewards with the Yale NLP lab and developing RL environments with Prime Intellect. Previously, I've done research in geometric algebra (APOLLO Labs), algebraic topology (SUMaC '23), abstract algebra (SUMaC '22), and biostatistics (Emory).
In my free time, I enjoy drawing, drinking coffee or tea, and exploring other ML and math disciplines.
Posts
Triton vs. Cutlass vs. ThunderKittens in Flash Attention 3
Coming soon! We develop Triton, Cutlass, and TK kernel implementations of Flash Attention 3 from scratch and test on H100s.
Whirlwind of PPO and RLHF for LLMs from scratch
RLHF with PPO from scratch and lots of fine-tuning GPT-2 models for movie sentiment classification. Transformer environments, adapative KL control, logit/temperature scaling, whitening, and more. Full implementation here.