Haelixe — A bare-metal deep learning engine in Rust

Architecture

Modern primitives, implemented from the ground up.

The exact mathematical foundations used by state-of-the-art architectures like LLaMA 3 and Mistral — written directly against Haelixe's own tensor engine, not borrowed from another framework.

01 / Autograd

A dynamic computation graph, tracked and torn down deterministically.

Every operation builds a directed acyclic graph on the fly. Gradients flow backward through reverse-mode differentiation, accumulated in-place via topological sort.

Graph Reverse-mode DAG, built dynamically per forward pass
Accumulation HashMap-backed, in-place gradient mutation

02 / Mixed precision

BF16 storage with an F32 stability boundary.

Model weights live in BFloat16 to halve VRAM footprint and PCIe bandwidth. At the compute boundary, a JIT autocast upcasts to F32 for numerical stability on hardware without native 16-bit ALUs.

Storage BF16 weights, halved memory and bandwidth
Optimizer Master-weights pattern: F32 state, BF16 model

03 / Tensor engine

Views, transposes, and slices touch zero bytes of memory.

Tensors are backed by one physical buffer. Reshaping, transposing, and slicing only rewrite the shape and stride arrays — the underlying data never moves. Broadcasting is implemented the same way, with zero-stride dimensions.

Layout Strided views over a single physical buffer
Broadcasting Zero-cost, via zero-stride dimensions

Systems engineering

Hardware-aware, on purpose.

Haelixe treats the memory allocator and the GPU dispatch layer as first-class parts of the framework — not an afterthought bolted on for performance later.

Binning slab allocator

Similar tensor sizes are grouped into shared pools, using power-of-two binning to raise cache-hit rates and cut driver-level VRAM overhead.

Deterministic RAII reclamation

Arc reference counting guarantees VRAM slabs return to the free list only once the final autograd reference is destroyed — no silent gradient corruption.

Kernel fusion & flash attention

The Linear layer fuses MatMul and bias-add into a single WGSL shader. Attention scores are computed entirely within GPU L1 cache via fused kernels.

Out-of-core data loading

Binary dataset files are mapped directly into virtual memory with memmap2, enabling training on datasets larger than system RAM.

In practice

Sequence denoising, end to end.

The reference implementation in haelixe-lab trains a Transformer block — RoPE, RMSNorm, GELU, cosine-annealed AdamW — to recover a clean signal from a noisy multi-frequency sine wave.

haelixe-lab/src/main.rs

use haelixe::{DType, Device, Shape, Tensor, TransformerBlock, RMSNorm, optim::AdamW};
use std::f32::consts::PI;

fn main() {
    let gpu = Device::gpu();
    let batch_size = 4;
    let seq_len = 32;
    let hidden_dim = 64;
    let num_heads = 4;

    let mut embed = haelixe::Linear::new(1, hidden_dim);
    let mut block = TransformerBlock::new(hidden_dim, num_heads);
    let mut final_norm = RMSNorm::new(hidden_dim);
    let mut head = haelixe::Linear::new(hidden_dim, 1);

    let mut optimizer = AdamW::new(0.001);
    let total_epochs = 100;

    for epoch in 0..total_epochs {
        let cosine_decay = 0.5 * (1.0 + (PI * epoch as f32 / total_epochs as f32).cos());
        // forward pass, backward pass, optimizer step…
        // converges from MSE 3.74 to < 0.85
    }
}

Clone it. Run it. Break it.

Haelixe ships as a Cargo workspace. haelixe-lab is the downstream consumer project — it verifies mathematical convergence and hardware integration on every change.

git clone https://github.com/sethigris/haelixe.git