GitHub - timothygao8710/minLlama: Yet Another Llama 3.2 implementation (in pure numpy)

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.py

Edit prompt, sampling_temperature, and the token budget at the top of each file to change generation.