A CUDA deep learning framework written from scratch in C++17, CUDA and Python:
tensor engine, reverse-mode autograd, hand-written kernels, an nn API, a
transformer stack, an NVRTC kernel-fusion compiler and a live GPU dashboard.
It trains real models. Every op is verified against numeric gradients, and the test suites require the optimizers to converge and a GPT to overfit a batch before they pass.
import numpy as np
import novatorch as nt
model = nt.nn.Sequential(nt.Linear(784, 256), nt.ReLU(), nt.Linear(256, 10))
opt = nt.AdamW(model.parameters(), lr=1e-3)
loss_fn = nt.CrossEntropyLoss()
logits = model(x) # x: [batch, 784] on the GPU
loss = loss_fn(logits, y) # y: int32 class ids
opt.zero_grad()
loss.backward()
opt.step()| Area | Details |
|---|---|
| Tensor engine | strided views (transpose/slice/permute/expand cost nothing), full NumPy broadcasting, float32 compute with int32 indices and fp16/bf16 storage |
| Autograd | reverse-mode graph built during the forward pass, topological execution, leaf accumulators, no_grad, nt.grad(...) |
| Kernels | 128x128 register-tiled GEMM (~2.7 TFLOP/s fp32 on an RTX 4050 laptop), fused elementwise + activation gradients, warp-shuffle reductions, row-wise softmax, im2col conv, fused optimizer steps |
| Attention | FlashAttention-2 style forward kernel (4.2 MB vs 138 MB live memory at seq=1024) plus a differentiable composed path used for training |
| nn API | Linear, Conv2d (groups, stride, padding, dilation), MaxPool/AvgPool/AdaptiveAvgPool, LayerNorm, BatchNorm2d, RMSNorm, Embedding, Dropout, 9 activations, 5 losses |
| Optimizers | SGD (momentum/Nesterov), Adam, AdamW, RMSprop, AMSGrad, gradient clipping, 7 LR schedulers |
| Transformers | multi-head attention with fused QKV, pre-norm blocks, GPT with training and sampling |
| JIT | expression -> CUDA C -> NVRTC -> cubin, cached; 1.8x on a fused relu(x*y+b) |
| Tooling | caching allocator, profiler, live dashboard, DataLoader with prefetch, MNIST/CIFAR-10, .npz checkpoints |
pip install pybind11 cmake numpy
pip install -e .Needs an NVIDIA GPU with compute capability 7.5+ and CUDA 11.8+. Details and troubleshooting: docs/building.md.
Verify:
python tests/test_framework.pypython examples/mnist_cnn.py --epochs 3 # CNN on MNIST
python examples/gpt_lm.py --steps 500 # character-level GPT
python examples/benchmark.py # kernel benchmarks
python examples/jit_demo.py # runtime kernel fusionMeasured on an RTX 4050 laptop (20 SMs, 6 GB):
MNIST CNN 12,600 images/s 94% test accuracy after 1 epoch
GPT (4L, 128d) 50,000 tokens/s loss 3.52 -> 0.05 in 200 steps
matmul 2048^3 6.2 ms 2,757 GFLOP/s
elementwise add 1.27 ms (2^24) 158 GB/s
training step 3.4 ms 74,900 samples/s (MLP 1024x3, batch 256)
core/ tensor engine, differentiable ops -> core/README.md
autograd/ backward nodes and the engine -> autograd/README.md
cuda/ allocator, streams -> cuda/README.md
cuda/kernels/ every CUDA kernel -> cuda/kernels/README.md
nn/ Module, layers, losses, optimizers -> nn/README.md
transformer/ attention, GPT -> transformer/README.md
compiler/ NVRTC JIT -> compiler/README.md
profiler/ kernel timing, metrics -> profiler/README.md
distributed/ NCCL multi-GPU (Linux) -> distributed/README.md
bindings/ pybind11 module -> bindings/README.md
novatorch/ Python package -> novatorch/README.md
tests/ C++ and Python suites -> tests/README.md
examples/ runnable programs -> examples/README.md
dashboard/ live GPU dashboard -> dashboard/README.md
docs/ quickstart, API, architecture, building -> docs/README.md
Every directory has a README explaining what it contains and what each file does.
python tests/test_framework.py # 25 groups
cmake --build build --config Release --target test_tensor
./build/Release/test_tensor.exe # 49 checksBoth suites gradient-check every differentiable op against central differences, and the Python suite additionally requires each optimizer to fit a known linear function, a classifier to reach >98% accuracy, and a small GPT to overfit a single batch.
Deliberate boundaries, so nothing here is a surprise:
- Compute is CUDA float32 (int32 for indices and labels). CPU tensors exist for I/O only.
- First-order gradients only: the backward pass is not itself recorded, so there is no double backward.
- No mixed-precision autocast and no in-place autograd ops.
- The FlashAttention kernel is forward-only (training uses the differentiable
composed path) and supports
head_dim <= 64. distributed/needs NCCL, so it is Linux-only and is not covered by the test runs on the development machine.
MIT. See LICENSE.