How dissecting Gemma 4 E2B-it's 20 KV-shared layers cut KV-cache memory by 57.1%, eliminated 432 XLA recompilations via device-tensor indexing in PyTorch (161.9 tok/s, 51.1x) and pure JAX lax.scan +...
./read_more.shExplore my thoughts on Machine Learning, AI, and Systems Engineering
Looking back at where it all started: the first meetup, the milestones, the challenges, and how building a developer community shaped my perspective.
Keshan Sanjaya Sodimana
Hands-on: train a ~10M-parameter GPT end-to-end on a free Colab/Kaggle TPU with JAX + Flax NNX, using the tpu-llm-from-scratch repo. Real commands, honest gotchas,...
Keshan Sanjaya Sodimana
A git-backed, hypothesis-driven optimization methodology that turns kernel tuning from art into science.
Keshan Sanjaya Sodimana
Deep-dive into TPU profiling - XProf, HLO, and the Roofline Model
Keshan Sanjaya Sodimana
Building a fused FlashAttention kernel on TPUs with JAX Pallas — from understanding the hardware to writing a working kernel that never materializes the...
Keshan Sanjaya Sodimana
A complete, beginner-friendly guide to processing large datasets in parallel using GCP Batch, Docker, and Cloud Storage from zero to a running job.
Keshan Sanjaya Sodimana
This guide demonstrates how to build a secure and reliable Python pipeline for video analysis using the Google Gemini API.
Keshan Sanjaya Sodimana