I wanted to understand what actually makes DeepSeek V3 different, so I rebuilt its three main ideas from scratch in pure JAX: Multi-head Latent Attention, mixture-of-experts with bias-based load balancing instead of an auxiliary loss, and multi-token prediction. The model is about 190 lines.

It’s a 36.2M-parameter version trained on a MacBook Air M5 (24 GB) with jax-metal, on 207 old detective novels. 12,000 steps took about 2.6 hours.

Full write-up with diagrams and results: https://vaibhaav.in/writing/deepseek-v3-jax/

The trained checkpoint and exact data are in the v1.0 release, so you can generate text without training.

Would love corrections if I got anything wrong :)