LucilleAI
A machine-learning framework written in Rust that keeps tensors on the GPU and records every operation for training
A machine-learning framework written from scratch in Rust, with a CUDA backend, that trains a real model.
What it is
I wanted to understand what actually happens between a tensor operation and a gradient update, so I built the whole stack: tensor library, autodiff, a neural-network layer set, an optimizer, a tokenizer, and a training loop. Then I wrote a model on top of it and watched its loss go down.
That is the honest framing. The value is not “this competes with PyTorch” — it does not, and I am not claiming it does. The value is that every layer is one I can open and read, and that I hit the specific problems that frameworks paper over.
What it does
- Tensors live on the GPU, and automatic differentiation keeps a recording of every operation so gradients can be replayed backwards.
- cuBLAS runs on its own CUDA stream while my code runs on another. The single most valuable thing in this project: matmuls overlap with the kernels around them, so the GPU is never idle waiting for work I have not handed it yet.
- The framework did not crash — it produced wrong numbers. A trace revealed a “ghost NaN” that only
appeared on Turing-generation hardware, where
tf32silently rounds tensor-core inputs. That single bug cost more time than everything else in the project combined, and the fix is one that only someone reading the generated code would ever find. - The tokenizer had a reproducibility bug. It counted tokens into a floating-point accumulator that cannot
reorder additions, so the same text produced a different vocabulary size on different runs.
u64fixed it. - A real model trained on real text: 3.7 million parameters, 128-token context, loss driven down from initialisation, and a tokenizer that round-trips Persian text without mangling it.
The numbers
| 81,200 | lines — 76,700 Rust, 4,500 CUDA |
| 24 | hand-written CUDA kernels |
| 32 | differentiable operations |
| 3.7 M | model parameters |
| 396 | pretraining steps, loss 0.0486 |
| 99.2% | accuracy on a held-out suite (656 of 661) |
Why it’s worth a look
Two reasons.
The first is the stream bridge. Overlapping library work with your own work is a real technique and it is genuinely hard to get right — you have to own memory lifetimes, because an allocation that is still in use by the other stream cannot be reused. The second is the NaN. I did not find it by reading code; I found it by dumping per-kernel traces and looking at where a finite number became not-a-number. It only happened on Turing, only under specific tensor-core paths, and every error-based check in the world would have passed it.
Current state
- AMP was off in both runs. No mixed-precision training was actually exercised.
- The project’s own evaluation flags itself as overfit, and it should. A 99% figure from a self-authored benchmark suite means the benchmark was too easy, not that the model is excellent.
- The model is small — 3.7 million parameters is a demonstration, not a language model.
None of that makes it less interesting. It makes it what it is: a complete, readable, from-scratch ML stack where every number on this page came from my own code and my own runs.
