A minimal, fully self-contained implementation of Llama 3.2 Instruct inference in ~100 lines of pure Numpy.
| File | Backend | Notes |
|---|---|---|
main.py |
Numpy | Naive, every step does prefill |
main_kv.py |
Numpy | Adds KV cache, $O(L^3) \rightarrow O(L^2)$ |
pytorch_llama.py |
Pytorch | Statically-shaped KV cache, streams only the relevant slice per token. Written to be easily hackable for research. |
jax_llama.py |
Jax | Streams the whole cache per token & does masking, for compilance with jax.jit
|
Ingredients in Llama 3.2: RoPE with the Llama-3 frequency scaling, GQA, RMSNorm, SwiGLU MLP, shared unembedding / embedding matrix.
Follow-up from https://github.com/timothygao8710/minWhisper.
Setup
Requires Python ≥ 3.10 and uv.
uv sync # Numpy only uv sync --extra torch # for pytorch_llama.py uv sync --extra jax # for jax_llama.py
Llama 3.2 is a gated model, so first request access on the model page, then authenticate and download the checkpoint (config, weights, tokenizer) into checkpoints/:
hf login hf download meta-llama/Llama-3.2-1B-Instruct \ --local-dir checkpoints/Llama-3.2-1B-Instruct \ --include "model.safetensors" "config.json" "tokenizer.json"
Usage
uv run main.py # or main_kv.py / pytorch_llama.py / jax_llama.pyEdit prompt, sampling_temperature, and the token budget at the top of each file to change generation.