An implementation-first course for building a small decoder-only transformer in MLX and proving it with logits, memory traces, and benchmarks.
Learning Outcomes
Implement RMSNorm, RoPE, SwiGLU, attention, and KV cache in MLX
Assemble a small decoder-only transformer path
Load weights and compare logits against a trusted reference
Measure prefill, decode, memory, and root-cause performance bottlenecks
Source-First Materials
course.md
NanoTransformer in MLX
Course Contract
NanoTransformer in MLX is the implementation-first bridge between local inference intuition and real systems judgment. The learner builds a small decoder-only transformer in MLX, then proves it works with reference checks, memory traces, and generation tests.
The course is not a tour of MLX APIs. It is a controlled rebuild of the core inference path:
tensors and lazy evaluation
RMSNorm, RoPE, SwiGLU, and attention
KV cache and continuous batching primitives
weight loading and logit matching
quantized linear layers
reproducible latency and memory measurements
Audience
Experienced engineers who can read Python and model configs but still feel one layer away from the actual inference stack.
Prerequisites
Local AI Model Systems, or equivalent understanding of weights, KV cache, quantization, and MLX basics
Comfort reading model configs and tensor shapes
Apple Silicon machine recommended for the labs
Course Outcome
By the end, the learner has a repo-quality mini runtime that can run a tiny decoder path, compare logits against a reference, explain memory peaks, and produce a systems memo about where time and memory go.
Chapters
MLX Execution and Tensor Building Blocks
Attention, RoPE, and KV Cache
Model Assembly and Weight Loading
Benchmarking, Forensics, and Systems Judgment
Source-First Rule
The canonical output is the artifact set in artifact-map.md, not a quiz score. LearnHouse export is optional and intentionally thinner than this course contract.
artifact-map.md
Artifact Map
Primary Artifacts
model/rms_norm.py: reference RMSNorm implementation and tests
model/rope.py: rotary embedding implementation with shape notes
model/attention.py: causal attention implementation with reference comparison
model/kv_cache.py: append-only KV cache and simple paged-cache sketch
model/block.py: decoder block with RMSNorm, attention, and SwiGLU
model/nano_transformer.py: full decoder-only model assembly
scripts/convert_weights.py: safetensors loading and name mapping notes
reports/systems_judgment.md: final memo about latency, memory, and correctness
Evidence Standard
An artifact is accepted only when it includes evidence:
unit tests for tensor shapes and edge cases
numerical comparison against a trusted reference where possible
memory and latency measurements with prompt length recorded
a short note explaining observed deltas
Final Portfolio
The finished portfolio should let future you answer:
How does a transformer block really move data?
Which tensors scale with context length?
Where does MLX lazy evaluation change measurement?
What breaks when a config assumption is wrong?
What would I optimize first and why?
references.md
References
Primary Reading Order
MLX documentation: arrays, lazy evaluation, compilation, and memory.
MLX examples and mlx-lm model implementations.
A small Llama-family or TinyLlama-style model config.
RoPE implementation notes from a reference transformer library.
Safetensors format documentation.
llama.cpp or MLX quantization notes for contrast after the dense path works.
Source Reading Tasks
Trace how mlx-lm represents model config fields.
Trace how a reference implementation applies RoPE offsets during generation.
Compare append-only cache behavior with any rotating or paged cache implementation.
Identify where a high-level implementation hides allocation or evaluation boundaries.
Verification Sources
Use primary source code or docs for claims about:
MLX evaluation semantics
model config fields
RoPE scaling behavior
safetensors storage
quantization formats
Avoid treating blog posts as final authority for numerical or implementation details.
labs/kv-cache-lab.md
KV Cache Lab
Goal
Implement a minimal KV cache and learn exactly what grows during generation.
Task
Build two versions:
Append-only contiguous cache.
Page-shaped sketch that reserves blocks and maps logical positions to physical slots.
The paged version does not need production kernels. It must make allocation behavior visible.
Output
Create:
model/kv_cache.py
tests/test_kv_cache.py
reports/kv_cache_tradeoffs.md
Pass Condition
The tests should prove that appends, reads, trim behavior, and variable-length requests preserve token order.
labs/logit-matching-lab.md
Logit Matching Lab
Goal
Prove that your NanoTransformer is not just plausible. It must match a trusted reference path for a tiny model or a reduced toy checkpoint.
Task
Load the same tokenizer and model weights into the reference implementation and your MLX implementation.
Run a fixed prompt through both.
Compare logits at the final position.
Report max absolute error and top-10 token agreement.
Explain any mismatch before changing code.
Output
Create reports/logit_matching.md with:
model ID and revision
prompt text and token IDs
dtype and quantization state
max absolute error
top-10 agreement
suspected causes for remaining delta
Pass Condition
For the dense fp16 path, target a small enough delta that the next-token top choices match. If exact tolerance is not realistic because the reference uses different kernels or dtypes, document why with evidence.
labs/memory-trace-lab.md
Memory Trace Lab
Goal
Separate weight memory, KV cache growth, temporary buffers, and benchmark artifacts.
Task
Instrument memory at these points:
Before import.
After tokenizer load.
After weight load.
After prompt prefill.
Every 128 generated tokens.
After cleanup.
Record prompt length, generated length, dtype, cache settings, and whether mx.eval was called before measurement.
Output
Create:
reports/memory_trace.csv
reports/memory_trace.md
The markdown report should include a root-cause explanation for the largest memory jump.
Pass Condition
You can explain the memory curve without hand-waving. Unknown memory should be called out as unknown, not hidden.
rubrics/capstone.md
Capstone Rubric
Artifact
Build a minimal decoder-only MLX runtime and produce a systems judgment report.
Rubric
Correctness, 3 points: shape tests pass, generation path is coherent, reference-logit comparison is documented.
Memory reasoning, 2 points: weight, KV cache, and temporary memory are separated with measured evidence.
Runtime reasoning, 2 points: prefill and decode are measured separately, with mx.eval boundaries explained.
Code clarity, 2 points: model components are small, named by architecture concept, and testable.
Judgment, 1 point: final memo states the next optimization and what evidence would change that choice.
Non-Passing Patterns
A model that only works for one prompt without shape tests.
A benchmark that reports one tokens-per-second number without context length.
A memory estimate with no measured comparison.
A fix that changes output without explaining the root cause.
Modules
MLX Execution and Tensor Building Blocks
Build the primitive tensor operations and execution-model discipline needed before assembling a transformer.
1.1 MLX Array and Execution Model
Objective
Build a concrete mental model for MLX arrays, lazy evaluation, and measurement boundaries.
Root Confusion
Many local-inference mistakes come from timing graph construction or Python dispatch instead of evaluated tensor work.
Theory
MLX can defer work until evaluation is forced. That makes code pleasant, but it also means a benchmark is wrong unless the measurement boundary is explicit.
Track:
array shape and dtype
whether an operation returns a view or new storage
when mx.eval is called
when memory peaks occur
Hands-On Task
Write small matrix multiply and elementwise benchmarks. Measure cold run, warm run, and run with explicit mx.eval.
Artifact
reports/mlx_execution_notes.md
Checkpoint
Can you explain why moving one mx.eval changes the measured latency?
Failure Mode
Mistaking a lazy graph timing for real compute throughput.
1.2 RMSNorm, Embedding, and RoPE
Objective
Implement the first transformer primitives: embedding lookup, RMSNorm, and RoPE.
Root Confusion
Position encoding often feels like metadata until offset handling breaks generation.
Theory
RMSNorm rescales hidden states without mean-centering. RoPE rotates query and key dimensions by position-dependent frequencies. During generation, the offset must match the position of the token being processed.
Hands-On Task
Implement RMSNorm and RoPE in MLX. Compare output against a small NumPy or PyTorch reference for fixed inputs.
Artifact
model/rms_norm.py, model/rope.py, and tests/test_rope.py
Checkpoint
What changes between prompt prefill RoPE positions and one-token decode RoPE positions?
Failure Mode
Correct shapes with wrong offsets, causing generation drift.
1.3 SwiGLU and MLP Block
Objective
Implement the feed-forward path used by modern decoder blocks.
Root Confusion
The MLP is often treated as boring, but it can dominate parameter count and memory bandwidth.
Theory
SwiGLU uses a gated activation path. For many Llama-style models, the MLP contains large up, gate, and down projections. Quantization and memory layout choices show up strongly here.
Hands-On Task
Implement a SwiGLU MLP with explicit shape checks and a reference comparison.
Artifact
model/mlp.py and tests/test_mlp.py
Checkpoint
Which dimensions determine MLP weight memory?
Failure Mode
Swapping gate and up projections or silently accepting wrong intermediate size.
Attention, RoPE, and KV Cache
Implement attention and cache state directly enough to reason about correctness and memory.
2.1 Causal Attention from First Principles
Objective
Implement causal self-attention without relying on a high-level transformer block.
Root Confusion
Attention is not just a formula. Runtime behavior depends on masks, cache layout, head grouping, dtype, and evaluation strategy.
Theory
Causal attention computes query-key scores, applies a future-token mask, normalizes with softmax, and combines values. Grouped-query attention changes the number of KV heads relative to query heads.
Hands-On Task
Implement a simple attention path for prefill. Add shape checks for query heads, KV heads, and head dimension.
Artifact
model/attention.py and tests/test_attention_shapes.py
Checkpoint
How does KV head count change cache memory?
Failure Mode
Broadcasting KV heads incorrectly while tests only cover one head configuration.
2.2 KV Cache Append and Read
Objective
Build the minimal KV cache needed for generation.
Root Confusion
The cache is not a vague optimization. It is persistent model state that must match committed tokens.
Theory
During prefill, the model writes keys and values for prompt tokens. During decode, it appends one or more new positions and reads all relevant prior positions.
Hands-On Task
Implement an append-only KV cache. Add tests for prompt prefill, single-token decode, and trimming.
Artifact
model/kv_cache.py and tests/test_kv_cache.py
Checkpoint
What exactly is stored per token per layer?
Failure Mode
Cache offset off by one, producing plausible but wrong generations.
2.3 Continuous Batching Skeleton
Objective
Build a tiny scheduler that exposes prefill/decode asymmetry.
Root Confusion
Batching is not just stacking prompts. Requests have different prompt lengths, decode lengths, and cache states.
Theory
Serving runtimes split prefill and decode because they stress hardware differently. Continuous batching lets new work enter while older requests continue decoding.
Hands-On Task
Write a scheduler simulation that handles variable prompt lengths and tracks per-request cache length.
Artifact
engine/scheduler.py and reports/scheduler_trace.md
Checkpoint
Why can a batch with many short requests behave differently from one long interactive request?
Failure Mode
Optimizing aggregate throughput while making single-user latency worse.
Model Assembly and Weight Loading
Assemble the decoder path, load weights, and prove correctness with reference comparisons.
3.1 Assemble Decoder Block
Objective
Assemble RMSNorm, attention, and MLP into a decoder block.
Root Confusion
Small residual-order mistakes can preserve shapes while changing model behavior completely.
Theory
A decoder block applies normalization, attention, residual addition, another normalization, MLP, and another residual addition. Exact ordering matters.
Hands-On Task
Implement one decoder block with explicit intermediate names and tests for deterministic toy weights.
Artifact
model/block.py and tests/test_block.py
Checkpoint
Where are the two residual connections?
Failure Mode
Passing shape tests while failing logit matching because residual order is wrong.
3.2 Load Safetensors and Map Weights
Objective
Load model weights and map external names to your local implementation.
Root Confusion
Weight loading failures are often architecture-name mismatches disguised as file or dtype problems.
Theory
Safetensors stores named tensors. A local runtime needs a mapping from checkpoint names to implementation names, plus validation that shapes match expected config-derived shapes.
Hands-On Task
Write a loader that reads a tiny checkpoint or subset and rejects mismatched shapes with useful errors.
Artifact
scripts/convert_weights.py and reports/weight_map.md
Checkpoint
Which config fields should be checked before assigning a tensor?
Failure Mode
Transposed or misnamed weights that load without error.
3.3 Logit Matching and Generation
Objective
Prove that the assembled model behaves like the reference path.
Root Confusion
Readable code is not evidence that the model is correct.
Theory
Logit matching catches architecture, dtype, RoPE, residual, and weight-mapping errors. Generation tests alone are too weak because many wrong implementations still emit plausible text.
Hands-On Task
Run the logit matching lab and then generate a short sequence from the matched path.
Artifact
scripts/check_logits.py and reports/logit_matching.md
Checkpoint
Why is top-token agreement weaker evidence than direct logit comparison?
Failure Mode
Debugging output quality before proving the forward pass.
Benchmarking, Forensics, and Systems Judgment
Measure quantization, prefill, decode, and memory before choosing an optimization.
4.1 Quantized Linear First Pass
Objective
Implement a first-pass quantized linear layer and understand what it does not solve.
Root Confusion
Quantization reduces stored weight size, but it does not automatically make every part of inference faster or smaller.
Theory
A quantized linear layer stores compact weights plus scales and possibly zero points. Runtime must dequantize or use kernels that consume quantized data directly.
Hands-On Task
Implement a simple group-wise int8 or int4 linear path and compare memory, latency, and output error against dense weights.
Artifact
model/quant_linear.py and reports/quant_linear_tradeoffs.md
Checkpoint
Where did memory improve, and where did latency not improve?
Failure Mode
Reporting smaller weights without measuring total runtime memory.
4.2 Benchmark Harness
Objective
Build a benchmark harness that separates prefill, decode, and memory.
Root Confusion
One tokens-per-second number hides the workload.
Theory
Prefill is prompt processing. Decode is new-token generation. They have different bottlenecks and should be measured separately with warmup and fixed prompts.
Hands-On Task
Write a benchmark runner that records prompt tokens, generated tokens, time to first token, decode tokens per second, peak memory, dtype, and cache mode.
Artifact
scripts/bench_decode.py and reports/benchmark_manifest.md
Checkpoint
Which metric would change first if KV cache memory became the bottleneck?
Failure Mode
Comparing runs with different prompt lengths or warmup states.
4.3 Systems Judgment Report
Objective
Turn implementation evidence into a decision memo.
Root Confusion
Optimization without a root-cause memo becomes random tuning.
Theory
Systems judgment means naming the bottleneck, showing evidence, proposing the next change, and stating what evidence would change your mind.
Hands-On Task
Write the final systems judgment report using your logit, memory, and benchmark artifacts.
Artifact
reports/systems_judgment.md
Checkpoint
What would you optimize first, and what measurement supports that decision?
Failure Mode
Choosing an optimization because it sounds advanced rather than because evidence points there.