Aether-ML

A deep learning framework built from the ground up based on two numerical libraries NumPy and CuPy, with hand written GPU kernels for both CUDA and ROCm.

Why Aether-ML?

Aether-ML implements forward and backward propagation, convolutional and fully connected networks, normalization, pooling, and losses without building directly off any existing deep learning libraries. NumPy provides powerful ndarray datatypes that vectorize operations; CuPy ports this existing logic to the GPU, giving the extra functionality of custom GPU kernels speeding up an already embarrassingly parallel task.

The GPU path is NOT a thin wrapper over vectorized array operations, and it does not rely on an xp alias for a unified layer pass. Instead, nearly all components each run as hand-written kernels compiled from source templates targeting both NVIDIA and AMD hardware. To combat CPU overhead, all fronts including device backend, kernel-variant selection, and layer building are resolved ahead of time at model construction, and never re-evaluated on every call.

Aether-ML is an intentionally ground-up project born knowing nothing about machine-learning nor GPU kernel programming. However, rather than promote a weak mental model of machine learning and how deep learning frameworks work, the goal was to implement how deep learning runtimes work under the hood. With this approach, all system level decisions, hardware dispatch, and memory footprints had to be resolved, building real knowledge about architecting and maintaining real, production-style software.

Key Features

Layers — Conv2d, Dense, BatchNorm, MaxPool2d, AvgPool2d, GlobalAvgPool, Dropout, SpatialDropout, ReLU, LeakyReLU, SoftMax, with He initialization (He et al.) and Philox counter-based PRNG for dropout masks (Salmon et al.)

Kernels — Hand-written CUDA and HIP kernels for every compute layer, magic-number integer division (Granlund & Montgomery) for on-device index decomposition, and numerical stability guards on optimizer moment denominators, standard scaling, and loss terms

Kernel Fusion — Automatic graph-level operator fusion on CuPy combining BatchNorm and ReLU/LeakyReLU on both forward and backward passes.

Training — AdamW with decoupled weight decay (Loshchilov & Hutter), L1 and L2 regularization executed inside the optimizer kernel, label smoothing within categorical cross entropy (Guo et al.), batch shuffling, in-training validation, and throttled progress reporting, and mixed precision via set_precision, which casts most layers while forcing FP32 in the optimizer, Softmax, and BatchNorm for numerical stability.

Serialization — .aether bundles pair a JSON manifest with safetensors weights (zero pickle). The manifest records the layer graph, the loss, optimizer, and accuracy configs, the precision policy, the seed, and the fitted preprocessing pipeline, so Model.load returns a finalized model that accepts raw inputs straight away. Loading ever resolves for Aether's own classes, holds unexpected zip bombs, or carrying weights where shape, name, or dtype fail when they don't match the manifest.

Preprocessing — Compose pipelines with Rescale, StandardScaler, and ToTensor

Installation

Aether-ML requires python 3.12 or newer. It depends on NumPy and safetensors and runs on the CPU out of the box. For GPU accelerated workloads via CuPy, there are separate wheels per vendor and toolkit versions, so it's recommended to consult the CuPy installation guide for supported acceleration.

Basic (CPU)

pip install aether-ml

With GPU Acceleration

Install a pre-configured target environment using extras. Pick the one matching your installed CUDA toolkit.

# NVIDIA CUDA 13  (gpu-nvidia is an alias for this)
pip install "aether-ml[gpu-nvidia-13]"

# NVIDIA CUDA 12
pip install "aether-ml[gpu-nvidia-12]"

AMD ROCm: Pre-built CuPy wheels depend on a specific ROCm driver version (e.g. pip install cupy-rocm-7-0). Refer to the CuPy ROCm installation guide for supported driver targets. For users with the ROCm 7.14.x version and a GPU that supports said version, the repository includes a VS Code dev container under .devcontainer/rocm-gfx1201/ for the AMD configuration.

Quickstart

A complete model lifecycle, from construction to a saved model.

import aether as ae
import numpy as np
import cupy as cp

TARGET_DEVICE="cupy"
feature_pipeline = ae.Compose([
    ae.ToTensor(dtype='float32', target_device=TARGET_DEVICE),
    ae.Rescale(factor=1.0 / 255.0),
    ae.StandardScaler()
])
# Setup
model = ae.Model()
model.manual_seed(seed=42) #applies for all compute and PRNG based layers

model.add(ae.Conv2d(3, 32, (3, 3), (1, 1), padding="same"))
model.add(ae.BatchNorm(epsilon=1e-5, momentum=0.9))
model.add(ae.ReLU())
model.add(ae.MaxPool2d((2, 2), (2, 2), padding="valid"))

model.add(ae.Conv2d(32, 64, (3, 3), (1, 1), padding="same"))
model.add(ae.BatchNorm(epsilon=1e-5, momentum=0.9))
model.add(ae.LeakyReLU())
model.add(ae.MaxPool2d((2, 2), (2, 2), padding="valid"))
model.add(ae.SpatialDropout(rate=0.1))

model.add(ae.Conv2d(64, 128, (3, 3), (1, 1), padding="same"))
model.add(ae.BatchNorm(epsilon=1e-5, momentum=0.9))
model.add(ae.ReLU())
model.add(ae.MaxPool2d((2, 2), (2, 2), padding="valid"))
model.add(ae.SpatialDropout(rate=0.1))

model.add(ae.GlobalAvgPool())
model.add(ae.Dense(128, 128, l2=2e-5))
model.add(ae.Dropout(rate=0.075))

model.add(ae.Dense(128, 10))

model.configure(
    loss=ae.SoftmaxCategoricalCrossEntropy(label_smoothing=0.05),
    optimizer=ae.AdamW(lr=0.001, decay=1e-4, weight_decay=0.01),
    accuracy=ae.CategoricalAccuracy(),
    preprocessor=feature_pipeline   # Optional, see examples/cifar10/
)
model.to(TARGET_DEVICE)
model.finalize(input_shape=X_train.shape[1:], fuse_kernels=True)

model.train(X=X_train, 
    y=y_train, 
    epochs=25, 
    batch_size=128, 
    shuffle=True,
    print_every=100,
    validation_data=(X_test, y_test)
)

model.evaluate(X=X_test, y=y_test, batch_size=128)
model.save(filepath="saved_models/quickstart.aether")

# --- In a separate evaluation / deployment script ---
model = ae.Model.load("saved_models/quickstart.aether")
predictions = model.predict(X_raw)

The accuracy here should converge to 75-77% test accuracy on CIFAR-10 across 25 epochs. For end to end dataset pipelines, mixed-precision configurations, routing model to a fused kernel path, and proper saving and loading functionality, see the complete scripts inside examples/cifar10/

Architecture

Four decisions shape most of the codebase.

Resolution happens once, not per call. model.to() triggers _compile_for_device on every component, swapping method pointers to backend-specific implementations and selecting kernel variants once, ahead of time. The cost is that layers have two lifecycle phases; an uncompiled construction phase layer.build() and a compiled execution phase.

One exception is deliberate. The training flag stays a plain runtime boolean, because resolving it ahead of time would double the bound variants in every layer to save tens of nanoseconds per layer.

One kernel source, two vendors. cupy.RawKernels in the framework are generated from a shared template, with vendor substitution maps for CUDA and HIP supplying the divergent pieces: intrinsic names, launch geometry, matrix-core APIs. This does not affect kernel performance, but adds an extra step during compilation, and affects readability.

Kernels compile per shape, not per call. Convolution kernels are cached on shape metadata, with filter size and stride baked in as compile-time constants rather than passed as arguments, letting the compiler unroll inner loops and fold index arithmetic. The cost is a compile on first encounter of each new shape, and a cache that grows with shape diversity, which is acceptable when shapes stay stable across thousands of training steps.

Optional components are objects, not None. The Accuracy, Optimizer, Preprocessor, and TrainingProgress components all include an optional Null object variant. This removes if x is not None from the training loop, further reducing CPU overhead during the model lifecycle.

Benchmarks

All figures train the same CNN on CIFAR-10 with the same seed and fp16 precision policy, on an AMD RX 9070 XT and Intel Core Ultra 265K unless noted. The dispatch path differs, and with it Conv2d's precision: only the fused path can run Conv2d in fp16, so the other two compute it in fp32. The scripts and saved results are in benchmarks/.

Convergence

The solid lines are the desktop: NumPy on the CPU, then vectorized CuPy and fused kernels on the 9070 XT. The dashed line is the fused path on an RTX 3050 Mobile. The laptop GPU reaches a loss of 1.8 about 4× later than the 9070 XT (1.5 s vs 0.37 s), but still 10× sooner than the 9070 XT without custom kernels (15.4 s).

Throughput

Per training step: 4.85 ms on the fused kernel path vs 202 ms on vectorized CuPy on the same GPU, with 1,865 ms for NumPy on CPU. With the GPU paths, the kernel work shows a 42x speedup against identical hardware running array operations.

Transient VRAM

Transient VRAM footprint over two training steps. Unfused CuPy dispatches allocate independent temporary buffers in VRAM for intermediate steps, resulting in transient spikes up to 420 MiB. Custom fused kernels eliminate the need to create temporary buffers by executing multiple actions in the same kernel within registers and shared memory, reducing peak transient VRAM by 5.3× (80.0 MiB peak). The RTX 3050 Mobile trace has the same shape and peak, within 0.07 MiB, so it isn't plotted separately.

Building from Source

The test suite requires the source repository. To set up a development checkout:

git clone https://github.com/AlexanderSoftCode/Aether-ML.git
cd Aether-ML

python -m venv .venv
source .venv/bin/activate        # Windows: .venv\Scripts\activate

# Core install (NumPy, safetensors) — CPU path only
pip install -e .

# Or, for the GPU path, install with the extra matching your CUDA toolkit.
# This replaces the line above; pick one.
pip install -e ".[gpu-nvidia-13]"
# pip install -e ".[gpu-nvidia-12]"

Running the test-suite:

python3 -m unittest discover tests

A single line tests both the NumPy and CuPy backends. Some of the coverage is included below:

• Numerical gradient checks on every compute layer
• Unit tests for losses, optimizers, preprocessing transforms, and the test harness itself
• Integration tests covering full Model lifecycles including Compose pipelines and save/load round trips

Limitations and Scope

• No NCHW input shapes. Tight coupling between the NHWC layout and the kernels was required for maximum throughput. The alternative would've been to support both tensor shapes, and find workarounds for the non-contiguous array views that NCHW brings.
• No autograd. Every backward pass is written by hand. This is not an omission, and outside the scope of this project.
• Single GPU. No data or model parallelism.
• No learning rate scheduling, early stopping, or checkpointing.
• No gradient clipping. The numerical guards in the optimizer and loss protect against division by zero; they do not bound gradient magnitude.
• Two accuracy metrics: categorical and regression.
• Fused kernel selection is automatic. When CuPy is present, the fused path is used; there is no runtime switch to force the vectorized path.
• CI covers the CPU path only. GitHub Actions runners have no GPU, so the CuPy and kernel tests are skipped there and verified locally.
• Only two tested configurations. An RX 9070 XT on ROCm, and an RTX 3050 Mobile on CUDA.

Future Work

• Transformer Block — Implementation of a general transformer block and ViT transformer block into the framework, with the fallback vectorized path as well as a fused kernel path.
• Data augmentation pipeline — passed into configure, applied per mini-batch (Gaussian blur, flips, rotations, inversion)

Origins

This project began as an implementation of Neural Networks from Scratch in Python by Harrison Kinsley and Daniel Kukieła, and has since diverged considerably from it. Three breaks are worth naming.

Data flow. The book mutates a self.output attribute on each layer. Aether-ML returns its outputs instead, which saves VRAM in layers such as Dense and BatchNorm, while enabling construction-time dispatch binding.

Backends. The book is NumPy only. Aether-ML dispatches to hand-written kernels across two GPU vendors from a shared template.

Scope. Convolution, pooling, and GPU execution do not appear in the book. Those are original implementations here, along with the kernel caching, vendor templating, and layer build phase that support them.

The pre-refactor implementations are preserved under examples/legacy/ for anyone who wants to see the first era of this project.

References

• NNFS Book: Kinsley, H., & Kukieła, D. (2020). Neural Networks from Scratch in Python.

• Adam Optimizer: Kingma, D. P., & Ba, J. (2014). Adam: A Method for Stochastic Optimization. arXiv preprint arXiv:1412.6980.

• AdamW Optimizer: Loshchilov, I., & Hutter, F. (2017). Decoupled Weight Decay Regularization. arXiv preprint arXiv:1711.05101.

• Batch Normalization: Ioffe, S., & Szegedy, C. (2015). Batch Normalization: Accelerating Deep Network Training by Reducing Internal Covariate Shift. In Proceedings of the 32nd International Conference on Machine Learning (ICML) (pp. 448-456).

• He Initialization: He, K., Zhang, X., Ren, S., & Sun, J. (2015). Delving Deep into Rectifiers: Surpassing Human-Level Performance on ImageNet Classification. In Proceedings of the IEEE International Conference on Computer Vision (ICCV) (pp. 1026-1034).

• Label Smoothing: Guo, C., Pleiss, G., Sun, Y., & Weinberger, K. Q. (2017). On Calibration of Modern Neural Networks. In Proceedings of the 34th International Conference on Machine Learning (ICML) (pp. 1321-1330).

• Philox PRNG: Salmon, J. K., Moraes, M. A., Dror, R. O., & Shaw, D. E. (2011). Parallel Random Numbers: As Easy as 1, 2, 3. In Proceedings of International Conference for High Performance Computing, Networking, Storage and Analysis (SC) (pp. 1-12).

• Integer Division via Multiplication: Granlund, T., & Montgomery, P. L. (1994). Division by Invariant Integers Using Multiplication. In Proceedings of the ACM SIGPLAN Conference on Programming Language Design and Implementation (PLDI) (pp. 61-72).

• CIFAR-10 Dataset: Krizhevsky, A. (2009). Learning Multiple Layers of Features from Tiny Images (Tech. Rep. TR-2009). Department of Computer Science, University of Toronto.

License

MIT License (LICENSE)

The implementation preserved under examples/legacy/ descends from the MIT-licensed Python code accompanying Neural Networks from Scratch in Python and carries its own LICENSE retaining the original copyright notice.