What it does
A set of progressive exercises and benchmarks across three accelerator stacks, plus two cross-stack projects. Each exercise has a sibling .out file with real captured stdout from the listed hardware.
- CUDA (Tesla T4): vector add, SAXPY bandwidth, shared-memory reduction, transpose/coalescing, tiled matmul
- TPU/JAX (v5e): jit + HLO, grad/linear regression, flax.nnx MNIST, vmap, shard_map
- MLX (Apple M4): array ops, linear regression, MLP MNIST, attention, tiny GPT decoder inference
How it works
Exercise (.py / .cu) → Run on target hardware → Captured .out → perf/ benches → Cross-stack projects
- Project A — softmax 4 ways: row-wise softmax on (4096, 4096) in MLX, JAX, raw CUDA and Triton; timings in
projects/A_softmax/bench.md. - Project B — matmul ladder + roofline: naive CUDA → tiled CUDA → Tensor Core cuBLAS → MLX → JAX MXU, tracking TFLOPS vs peak; see
projects/B_matmul/roofline.md. - Each stack has a
perf/directory: matmul size sweeps and memory/HBM bandwidth benches; projects haverun_all.shrunners with cachedresults.txt.
Tech stack
CUDA (nvcc, sm_75)TritonJAX / Flax / OptaxTPU v5eMLX / MetalApple M4