I ported Andrej Karpathy’s llm.c to Mojo, extending @dorjeduck’s llm.
(written for Mojo 25.5, CPU only) with GPU kernels, in honor of Mojo’s v1.0.0 release. It runs on nightly 1.0.0b3. @dorjeduck showed Mojo could match C on the CPU; I wanted to see whether it could hold its own on the GPU, including Apple Silicon. I was personally motivated based on feedback I received that my coding and math expertise weren’t strong enough, and writing the fundamentals of generative models from scratch, just like writing a compiler from scratch, is the first step in closing those gaps. It also took a lot longer then I planned: I was hoping to have it done in two weeks, and it ran closer to four.
What’s in it:
- From-scratch autograd in pure Mojo: forward pass and hand-derived backpropagation (the derivation is ~1,470 lines of LaTeX in
docs/backprop.tex). All kernel and trainer code is hand-written without LSPs or LLMs; tests and the later optimization campaign were AI-assisted, with formal disclosure in-repo. - Full GPT-2 124M training: AdamW, mixed precision (fp32/bf16), gradient accumulation, checkpointing with resume, dataloaders, tokenizer, ZeRO stages 1 to 3 implemented (full stage-3 verification is a roadmap item). GPT-3 hyperparameter tables ship in the trainer.
- Three backends: CPU, CUDA (cuBLAS/cuBLASLt fast paths), and Apple Metal, plus a vendor-neutral portable GPU path. The kernel library is 16,993 lines of Mojo across 26 modules.
- A benchmark harness with llm.c vendored as a submodule and matched hyperparameters across every arm.
Numbers. On an M4 Max (B=4, T=1024, cold GPU, 30 s inter-arm cooldowns). llm.c has no Metal port, so PyTorch MPS is the baseline there:
| configuration | mean ms/step | tok/s | vs PyTorch MPS |
|---|---|---|---|
| llm.mojo bf16 | 498.92 | 8210 | 1.72x faster (vs MPS bf16) |
| llm.mojo fp32 | 652.06 | 6282 | 1.27x faster (vs MPS fp32) |
| PyTorch MPS fp32 | 830.27 | 4933 | baseline |
| PyTorch MPS bf16 | 857.26 | 4778 | baseline |
On an NVIDIA GB10 (DGX Spark, same config), bf16 reached parity with llm.c CUDA: 135.97 vs 135.77 ms/step median (30,154 vs 30,210 tok/s), and ~3.8x faster than PyTorch CUDA bf16. Honest note: fp32 is 1.40x behind llm.c. bf16 is the shipped default config, but I am willing to take another pass at optimizing fp32. The first working Metal port was 3627 ms/step, so the final bf16 number is 7.3x faster than where I started.
Reproduce it:
git clone --recurse-submodules git@github.com:ulmentflam/llm.mojo.git
cd llm.mojo
make install # pixi env + hooks (make install-cuda on NVIDIA)
make data # TinyShakespeare + GPT-2 weights
make train
make benchmark-metal # all four Metal arms, with mandatory cooldowns
Correctness gates. test_gpt2.c is ported to test_gpt2.mojo: 16 gradient tensors checked against a PyTorch reference plus a 10-step loss trajectory asserted at every step (tolerance 0.01, matching llm.c). A 235-test pytest equivalence suite deliberately exercises odd shapes (channels=767, head_dim=5) to catch invariant-based fast paths.
Gotchas I hit (compiler and stdlib feedback welcome):
-
Casting a SHARED (threadgroup) pointer to the GENERIC address space is a no-op on CUDA, but on Metal AIR, GENERIC means device memory. The compiler accepts it silently and every load reads zeros. Attention computed softmax over all-zero scores: perfectly uniform attention weights, no NaN, no crash, a valid-looking model that learns nothing. Fix: a rebind that preserves the SHARED annotation. I’d love a warning here.
-
Metal’s DeviceContext is a single in-order command queue, so the
ctx.synchronize()fences that are load-bearing on CUDA (without them: CUDA_ERROR_ILLEGAL_ADDRESS) are pure waste on Metal. Stripping them behindcomptime if not HAS_METALsaved roughly 90 ms/step. -
Metal AIR limits bite at compile time: fully unrolling a 4x32 comptime loop nest exceeded the function-body limit and crashed the Metal compiler (“failed to compile metallib”; fix was making the inner KV loop a runtime loop), and the 32 KB threadgroup memory cap forced NVIDIA tile sizes down to Br=8/Bc=8.
The full catalog is in docs/ai/metal_port_gotchas_and_optimizations.md.
My takeaways on Mojo as a GPU language, after ~17k hand-written kernel lines:
- I went in assuming the compiler would do the heavy lifting on portability and efficiency across hardware. In practice it was less agnostic and less compiler-offloaded than I expected: I ended up branching logic and building sets of device-specific custom paths per vendor, well beyond a CPU/GPU split. A lot of those device-specific optimizations feel like they belong in either the higher level tiling abstraction or the compiler for Mojo to become a more efficient GPU programming language.
- The CPU is the clear winner today. It took very little heavy lifting to get 4.0x over llm.c with OpenMP at 20 threads.
- Asynchronicity is the thing I want most. Kernels like FlashAttention 4 require async to maximize Blackwell, and warp groups are fading as asynchrony and asymmetric pipelining becomes the dominant pattern. A first-class, ergonomic way to setup asymmetric pipelining now would set Mojo up well for where GPU development is heading.
- There is more scaffolding than I expected. Some of that is correctable on my side, but I still wrote a lot of similar blocks over and over where CUDA can be much more compact (Karpathy’s implementation is the reference point).
- Hand-writing the backward pass was the right call, and I’d add autograd carefully. Deriving the gradients by hand was easy enough, and every op in
llmm/has an explicit_bwd(matmul_bwd,layernorm_fused_residual_bwd,attention_bwd). The real payoff was not the derivation, it was what I noticed while writing each backward kernel. Which activations had to be saved versus recomputed, and which gradients had to be communicated versus kept local. That is what drove most of the wins (store-P over recompute, targeted gradient zeroing instead of a full-buffer memset, reduce-scatter instead of allreduce once ZeRO shards). An op-by-op autograd, would have a harder time finding those. So if Mojo adds autograd, the feature I’d want most is a first-class custom-backward escape hatch (PyTorch’sautograd.Function, JAX’scustom_vjp) so the hot kernels keep a hand-tuned backward while the rest is differentiated for me. Whole-graph capture that can fuse epilogues (bias+GELU) and overlap collectives with compute would be an even bigger win. Either way, a compiled backward pass is a genuinely nice thing to have. - The
commcollectives are the other half of a real training framework, and they should be portable. ZeRO here is built directly on MAX’scommpackage (comm.allreduce,comm.allgather,comm.reducescatter); paired with hand-written kernels, that was enough to stand up all of ZeRO stages 1 to 3. The catch is that those collectives are CUDA-only (they lean on NVLink and CUDA IPC P2P), so multi-GPU ZeRO raises on Apple GPUs and I fell back to hand-rolled CPU stand-ins. A portable path would help a lot: simulate the collectives on CPU and Metal for compatibility, or back them with Apple’s clustering the way MLX does now (its Ring backend, and JACCL over Thunderbolt 5). Then the ZeRO optimization would not have to be dropped on non-CUDA targets.
Roadmap: full ZeRO-3 verification, AMD GPU verification, then Mamba1/2/3 and MoE.
Feedback I’m looking for: Mojo GPU API ergonomics (especially address-space semantics on Metal and any roadmap thinking on asymmetric pipelining), whether the comptime vendor dispatch pattern (HAS_CUBLAS / HAS_METAL) is idiomatic, and holes in the benchmark methodology. Any other non-obvious gaps in implementation and overall methodology.
Thanks: Thanks to @dorjeduck for the original Mojo port and to Karpathy for llm.c.