Running DeepSeek‑V3 in ~200 Lines of Pure JAX – A Hands‑On Walkthrough

1. What was released / announced

A community member just dropped a Show HN post titled "DeepSeek V3 in ~200 lines of pure JAX" (GitHub: https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax). In a compact repository they re‑implemented the DeepSeek‑V3 architecture – a 7B‑parameter decoder‑only LLM – using only JAX and Flax, without any heavyweight libraries like 🤗 Transformers. The whole model loading, tokenisation, and inference pipeline fits into roughly 200 lines of code.

2. Why it matters

a. Pure JAX = portable & performant

JAX has become the de‑facto standard for research‑grade tensor ops because of its XLA‑backed compilation and seamless multi‑accelerator scaling. By keeping the implementation pure JAX we get:

• Zero‑copy GPU/TPU execution – the model lives in device memory from the get‑go.
• Easy sharding – pmap or pjit can split the 7B parameters across multiple devices with only a few extra lines.
• Future‑proofing – as JAX adds new primitives (e.g., named_sharding) the code can adopt them without a massive refactor.

b. Minimalism is a great learning surface

Most open‑source LLM repos bundle a mountain of utilities (tokenizers, checkpoint loaders, optimizer wrappers). Stripping everything down to the essentials lets engineers see exactly how a transformer works and where the performance knobs are. This is invaluable when you need to:

• Debug latency spikes in production.
• Port the model to a custom inference service (e.g., a FastAPI container on GKE).
• Experiment with quantisation or sparsity without fighting an opaque library.

c. Cost‑effective inference for startups

Running a 7B model on a single A100 costs roughly $0.30 per hour for inference. With JAX’s jit‑compiled forward pass you can shave 15‑20 % off that number compared to a naïve PyTorch loop. For a SaaS that serves a few thousand requests a day, those savings add up quickly.

3. How to use it

Below is a quick‑start guide that gets you from cloning the repo to serving a REST endpoint with FastAPI. I ran these steps on an Ubuntu 22.04 VM with an NVIDIA A100 and the latest CUDA toolkit.

Step 1 – Clone & install dependencies

# Clone the repo git clone https://github.com/Vaibhaav-Tiwari/deepseek-v3-jax.git cd deepseek-v3-jax # Create a virtualenv (optional but recommended) python3 -m venv .venv && source .venv/bin/activate # Install JAX with CUDA support and FastAPI pip install "jax[cuda12_pip]==0.4.23" flax==0.8.2 tokenizers fastapi uvicorn 

Tip: Replace cuda12_pip with the version that matches your driver. The repo pins JAX to 0.4.23; newer releases work as well.

Step 2 – Download the checkpoint

The author provides a script to pull the pre‑converted DeepSeek‑V3 checkpoint from HuggingFace:

python download_ckpt.py --repo deepseek-ai/deepseek-v3-7b # The script saves `deepseek_v3_7b.npz` in the `checkpoints/` folder. 

Step 3 – Load the model (pure JAX)

Create a tiny wrapper that JIT‑compiles the forward pass. This mirrors the model.py from the repo but adds a convenience function for inference.

import jax import jax.numpy as jnp from flax import linen as nn from model import DeepSeekV3 # <- the 200‑line implementation # Load parameters (the .npz file stores a flat dict) params = jnp.load('checkpoints/deepseek_v3_7b.npz', allow_pickle=True)['arr_0'].item() # JIT‑compiled generation step (single token) @jax.jit def next_token(params, token_ids, cache): logits, new_cache = DeepSeekV3().apply(params, token_ids, cache, mutable=['cache']) # Greedy decode for simplicity next_id = jnp.argmax(logits[:, -1, :], axis=-1) return next_id, new_cache 

Step 4 – Spin up a FastAPI endpoint

from fastapi import FastAPI, HTTPException from tokenizers import Tokenizer app = FastAPI() # Load a SentencePiece tokenizer that matches the checkpoint tokenizer = Tokenizer.from_file('tokenizer/tokenizer.json') @app.post('/generate') async def generate(prompt: str, max_new_tokens: int = 64): # Tokenise prompt ids = tokenizer.encode(prompt).ids token_ids = jnp.array([ids]) cache = {} output_ids = ids.copy() for _ in range(max_new_tokens): nxt, cache = next_token(params, token_ids, cache) nxt_id = int(nxt[0]) output_ids.append(nxt_id) token_ids = jnp.array([[nxt_id]]) # Stop on EOS token (usually id 2) if nxt_id == 2: break return {"completion": tokenizer.decode(output_ids)} if __name__ == '__main__': import uvicorn uvicorn.run(app, host='0.0.0.0', port=8000) 

Run the server:

uvicorn main:app --host 0.0.0.0 --port 8000 

You now have a self‑hosted DeepSeek‑V3 inference service that runs entirely on JAX. Because the model is pure‑JAX you can later replace the next_token function with a sharded version (pmap) to scale across 4‑8 GPUs without touching the FastAPI layer.

4. My take

Having spent the last two years building the inference stack for Griffin AI Tech, I’m constantly juggling three competing goals: latency, cost, and observability. Most of the time we reach for 🤗 Transformers + ONNX Runtime because the ecosystem is mature and the integration with Kubernetes is straightforward. However, the DeepSeek‑V3 JAX implementation reminded me why a bare‑metal approach still has a place.

Engineering trade‑offs

In practice, I would use the pure JAX version when:

1. I need to squeeze the last few percent of performance on a fixed GPU fleet.
2. I want to experiment with custom sharding or mixed‑precision pipelines.
3. My team is comfortable debugging XLA HLO traces.

Conversely, for a quick MVP or when the team is more product‑focused, the Transformers route wins.

Real‑world use case I’m eyeing

At Griffin we are building a knowledge‑base assistant for enterprise SaaS platforms. The workflow is:

1. Retrieve relevant documents (vector DB).
2. Prompt‑engineer a short context.
3. Run a 7B LLM to generate a concise answer.

Using the JAX implementation, we can host the model on a single A100 inside a GKE node pool, expose it via a gRPC service, and then scale the retrieval layer independently. Because the inference code is just a few functions, adding a LoRA adapter for domain‑specific fine‑tuning becomes a matter of swapping a small parameter matrix – no need to rebuild a massive Docker image.

Bottom line

The DeepSeek‑V3‑JAX repo isn’t a plug‑and‑play product; it’s a learning artifact that shows how far you can go with a minimal codebase when you trust the underlying compiler. For engineers who love to peek under the hood, it’s a goldmine. For production teams, it offers a pathway to a highly‑optimized, cloud‑native inference service that can be containerised, autoscaled, and monitored with the same tooling we already use for other JAX workloads.

Give it a spin, profile the latency, and decide if the extra engineering effort buys you the performance you need. As always, the best tool is the one that fits the problem – and now we have another option in the toolbox.

Happy hacking!

— Naveen Malothu, AI Infrastructure Engineer & Founder, Griffin AI Tech