# deepseek-v3-jax — DeepSeek V3 in 190 lines of pure JAX

> deepseek-v3-jax rebuilds DeepSeek V3's key ideas in 190 lines of plain JAX that you can read in one sitting.

- Org: Vaibhaav Tiwari
- Date: 2026-11-13
- Category: repo, tutorial
- Source: https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax
- Canonical page: https://ai-tldr.dev/releases/deepseek-v3-jax-190-lines/

## In short

deepseek-v3-jax is an MIT-licensed, 190-line JAX version of the DeepSeek V3 architecture. The 36.2M-parameter model trained in about 2.6 hours on a MacBook Air M5. The author calls it an architecture demo, not a strong language model.

## What it is

deepseek-v3-jax is a small, readable copy of the DeepSeek V3 design written in pure JAX, with no Flax, Equinox or PyTorch. It has Multi-head Latent Attention, a DeepSeekMoE layer with 1 shared and 16 routed experts (top-4), bias-based load balancing without an auxiliary loss, and multi-token prediction. The model has 36.2M parameters, 13.6M of them active per token.

## How it works

Parameters are plain dictionaries of arrays, and compute runs in bf16 on f32 master weights. The author trained for 12,000 steps (about 98M tokens from 207 books) in roughly 2.6 hours on a MacBook Air M5 with 24 GB, using the jax-metal plugin. A per-expert bias changes which experts the router picks, not their weights. At step 2000 each expert got 5.7–6.9% of tokens with the bias, against 0.09–18.6% without it.

## Why it matters

The write-up shows that balancing experts with a bias cost almost no accuracy: validation loss was 3.921 with it and 3.923 without. It also shows Multi-head Latent Attention caching 144 numbers per token per layer, against 480 for standard attention in this model. The author says the output reads like detective fiction but the stories do not hold together.

## Who it's for

ML engineers learning how DeepSeek V3 works

## Try it

git clone https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax.git && cd deepseek-v3-jax && uv venv -p 3.12 .venv

## Links

- [GitHub Repo](https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax)
- [Write-up](https://vaibhaav.in/writing/deepseek-v3-jax/)
- [Hacker News](https://news.ycombinator.com/item?id=50497589)