JAX Transformer + MoE Kernel Optimization
Jun 2026JAX · Flax · Optax · TPU · Pallas · MoE
I built a decoder-only Transformer in JAX, Flax, and Optax for arithmetic reasoning, then extended it with top-1 mixture-of-experts routing. From there, I wrote a fused Pallas projection kernel to reduce activation memory traffic and compared it with a ragged-dot baseline. The project follows the same workload from model design down to accelerator execution.
Speculative Decoding Runtime & GPU Kernel Lab
Apr 2026Python · PyTorch · Triton · CUDA · C++ · CUDA Graphs · Nsight Systems
I built an inference runtime for speculative and speculative speculative decoding with fused GPU kernels, CUDA graph replay, and execution overlapped across two GPUs. With Qwen2.5-7B in FP8 and a trained EAGLE-3 drafter, it reached 697 decode tokens/s on one request on two B300s, about 2× my tuned autoregressive baseline, and ran a single request to 128,000 tokens. Along the way I built an exact-match test suite, tracked down a KV-cache stride bug, and found new matmul schedules worth 17–21% end to end.
GPT-2 Inference Optimization
Oct 2025 - Dec 2025C++ · CUDA · Slurm · Nsight · WMMA · KV Cache · HPC
I profiled GPT-2 inference on an NVIDIA A40 to understand where GPU time was going, then focused on matrix multiplication. I explored shared-memory tiling, launch configuration, vectorized loads, and a WMMA path, and added KV caching for autoregressive generation. I compared kernel behavior with the full inference path to understand how local improvements carried through to the application.