DeepSeek V3 Rebuilt in 190 Lines of Pure JAX
Original: Show HN: DeepSeek V3 in ~200 lines of pure Jax
Why This Matters
Minimal reimplementations help researchers audit and understand frontier model architectures directly.
Developer Vaibhaav Tiwari reimplemented DeepSeek V3's full architecture — MLA, DeepSeekMoE, aux-loss-free balancing, MTP — in 190 lines of pure JAX, training the model from scratch on a MacBook Air.
Tiwari's open-source project strips DeepSeek V3 down to its architectural essentials without any high-level ML frameworks — no Flax, no Equinox, no PyTorch. All model parameters are plain dicts of arrays; every component is a readable top-to-bottom function in a single file, deepseek_v3.py.
The implementation covers Multi-head Latent Attention (MLA), which compresses the KV cache from 480 to 144 numbers per token per layer. Decoupled RoPE keeps a single rotary key shared across all heads outside the compressed latent. The DeepSeekMoE setup includes one shared expert that always runs alongside 16 routed experts. Aux-loss-free load balancing and Multi-Token Prediction (MTP) round out the core architecture features.
Training ran from scratch on public-domain detective fiction — a deliberately small corpus chosen to keep the experiment MacBook Air-friendly. The repo has 142 stars and 14 forks since posting, with a detailed write-up including diagrams available at vaibhaav.in. The project is MIT-licensed.