The Problem

Running a large‑language‑model on Apple Silicon is hard to learn from a black‑box library. Users need a small, end‑to‑end codebase that shows how tensors are turned into logits, how KV caches are built, and how requests are scheduled. Current tutorials either assume a CUDA backend or hide the math behind high‑level APIs, leaving systems engineers without a clear path to optimise memory and compute.

What This Does

tiny-llm provides a hands‑on, course‑style codebase that implements a Qwen‑3 model entirely with MLX arrays and a lightweight METAL kernel runtime.

  • Week‑1 core builds attention, RoPE, GQA, RMSNorm, the MLP, and the autoregressive loop in Python.
  • Week‑2 adds a request‑bounded KV cache, quantised mat‑vec kernels, SIMD prefill, and tiled prefill, exposing the full inference pipeline.
  • Week‑3 introduces continuous batching, chunked admission, and paged KV, turning the system into a mini vLLM.
  • Week‑4 adds a bounded, validated agent loop with tool‑driven steering.

The implementation lives in src/tiny_llm/ while the reference solution is in src/tiny_llm_ref/. Test coverage is provided in tests/ and tests_refsol/, and the book pages in book/ document the theory.

How It Is Wired

Entry point – main.py (in the root).

  • Calls pdm run check-installation for dependency sanity.
  • Imports tiny_llm.model (from src/tiny_llm/model.py) to load the Qwen‑3 weights.
  • model.load_weights() reads the checkpoint via model_names.py and stores tensors in MLX arrays.
  • The inference loop (tiny_llm.inference.autoregressive_loop) calls tiny_llm.attention.forward, tiny_llm.rmsnorm.forward, and the quantised mat‑vec kernel (src/extensions/quantized_matmul.cpp) compiled by src/extensions/CMakeLists.txt.
  • Batch scheduling is handled by tiny_llm.scheduler.schedule, which is exercised by batch-main.py.
  • Agent behaviour is exercised by agent.py, calling tiny_llm.agent.run.

Side‑effects – All writes are in-memory; the only filesystem touch is the checkpoint read and optional log files. No external network calls occur. Critical hub – src/tiny_llm/inference.py orchestrates the call chain and is the widest blast radius; a change here propagates to all weeks.

How To Use It

  1. Setup ``bash git clone https://github.com/moses-y/tiny-llm cd tiny-llm pdm install -v ``
  2. Configuration – No environment variables are required; the checkpoint is pulled automatically from the URL in model_names.py.
  3. Running ``bash pdm run test-refsol -- -- -k week_1 # verify week‑1 implementation python main.py --model qwen3-4b --prompt "Hello" ` For agent mode: python agent.py --prompt "Explain recursion"`

The CI (.github/workflows/main.yml) runs pdm run test-refsol on every push.

Real‑World Use

A dev team on macOS can iterate on KV‑cache bandwidth by modifying src/tiny_llm/kv_cache.py, re‑running bench.py from benches/, and comparing the JSON results in benchmark_results/. This workflow mirrors a production deployment: load weights, execute a batch, profile GPU use, and iterate.

Code Health & Issues

  • Low – Missing lockfile for pyproject.toml (non‑reproducible builds).
  • Low – 133 Python files but only 3 test suites in the root; many exercise files lack unit tests.
  • Low – No Dockerfile; the repo relies on local Apple Silicon.

No critical security or licensing gaps were found. The CI covers code style and coverage for reference solutions, but the student implementation is not fully tested.

The Bottom Line

tiny-llm is a focused, Apple‑silicon‑friendly sandbox that lets systems engineers see every step of LLM inference. It is ideal for educational use or rapid prototyping on a single Mac. The main limitation is the lack of a lockfile, which may surface dependency drift in longer‑term projects.