NEEDLE
A machine-learning library built from scratch with automatic differentiation, CNNs, LSTMs, and Transformers, accelerated on both NVIDIA and AMD GPUs.
NEEDLE is a machine-learning library I built entirely from scratch, one of the most challenging and rewarding projects I've taken on. It came out of a desire to understand and control every layer of the ML pipeline, from core tensor operations to the hardware acceleration that underpins performance. What started as a simple automatic-differentiation engine for tensor operations grew into a full, high-performance framework capable of training CNNs, LSTMs, and Transformers, an exercise in learning, debugging, and pushing what's possible with low-level custom code.
Note: This was a course project, so I keep the source private. Happy to walk through it or share on request.
Automatic differentiation
The project began with automatic differentiation (autodiff), the concept at the heart of backpropagation. Computing gradients efficiently is what makes training possible, so I implemented autodiff to derive arbitrary tensor operations, the foundation everything else is built on. This meant understanding how derivatives propagate through functions and how to store and track operations within a computational graph, then building a smooth, scalable way to perform those calculations without leaning on an existing framework. It gave me full control over gradient computation and set the stage for everything that followed.
The library: abstraction
With autodiff working, I turned to the abstraction layer, a custom, numpy-like library for tensor operations like addition, multiplication, and reshaping. The aim was something flexible, modular, and performant enough to eventually carry real ML workloads. This layer matters because it provides the building blocks for heavier operations like matrix multiplication, convolution, and optimization, so I designed the architecture carefully to stay both fast and easy to extend.
The library: modules
On top of the core tensor operations, I built more specialized modules, from matrix multiplication to convolution, which is essential for deep learning. With those in place I could train real models: Convolutional Neural Networks (CNNs), Long Short-Term Memory networks (LSTMs), and Transformers. The modular design kept the framework flexible, letting me add new algorithms over time while preserving the core principles of efficiency and simplicity. Each module was crafted to slot cleanly into the broader framework, making complex models straightforward to assemble from the ground up.
Hardware acceleration with CUDA
No matter how clean the code, without the GPU performance would always be the bottleneck, so a key goal was hardware acceleration. I integrated CUDA so tensor operations could run on the GPU, which meant mastering low-level memory management and parallelism. The payoff was substantial: training on the GPU dramatically cut the time to convergence, making it feasible to experiment with larger, more complex models and datasets. This was the step that took the project from a proof of concept to a genuinely usable tool for real machine-learning tasks.
AMD support with HIP
After CUDA, I wanted to support a broader range of hardware, since many practitioners run AMD GPUs. I added AMD support through HIP (the Heterogeneous-compute Interface for Portability), a cross-platform layer that lets the same code target both NVIDIA and AMD hardware. This required understanding the nuances of both CUDA and HIP as well as the underlying AMD architecture, but it let the library run efficiently across a much wider range of hardware configurations.
Reflection
Looking back, I'm proud of how far NEEDLE came, from a simple autodiff implementation to a full library that can train state-of-the-art models across different hardware platforms. The work demanded real investment in learning, debugging, and problem-solving, and in return it gave me a deep understanding of how ML frameworks actually work under the hood. Most of all, it reinforced the value of owning the entire stack, from the algorithmic foundation to hardware acceleration, and the freedom to customize every layer to fit the problem at hand.