Rustral

A Rust neural network framework for auditable, backend-agnostic NLP research

---

What Is Rustral?

Rustral is a 24-crate Rust workspace for building, training, and deploying neural networks. It is designed around three commitments that are unusual in the deep learning framework landscape:

1. No hidden global state. Every forward pass receives an explicit ForwardCtx carrying the backend, training/inference mode, run ID, optional shape policy, and optional operation profiler. There is no silent model.train() / model.eval() toggle, no global tensor registry, and no implicit mode switching.

2. Backend-independent model definitions. Layers are written against the Backend and TensorOps traits. The same model code runs unchanged on a reference CPU backend (ndarray with SIMD), an optimized Candle backend (CUDA/Metal), and an experimental cross-platform WGPU backend (native WGSL compute shaders).

3. Reproducibility by construction. Every benchmark, NLP training run, and LLM inference run emits a schema-validated JSON manifest recording machine metadata, git SHA, dataset checksums, hyperparameters, and raw timing distributions with 95% confidence intervals. These manifests are validated in CI.

Rustral is the subject of an EMNLP 2026 systems paper (docs/paper/emnlp_paper.tex) that presents the first systematic 6-framework cross-framework operator benchmark and proposes Cross-Backend Consistency Regularization as a training-time methodology for preventing silent implementation drift.

Who Is It For?

• Students learning how networks work under the hood
• Researchers who want reproducible, typed model code with auditable execution
• Engineers experimenting with Rust-native ML pipelines and deployment

---

Key Results at a Glance

6-Framework CPU Operator Comparison

Mean wall-clock ms on identical workloads (AMD EPYC 7702P, 64-core). Best per row in bold.

GPU Comparison (CUDA, RTX 2080 Ti)

On-Device Kernel Speedups (CPU SIMD)

NLP Reproducibility (deliberately minimal baselines)

LLM Inference

Low wrapper cost on one matched path: Rustral's transformer pipeline matches raw Candle at 1.09 ms in the reported workload. This is evidence for that measured path, not a general no-overhead claim.

Operator-to-task transfer: 12.8× LayerNorm speedup becomes 1.41× at the NLP pipeline level because normalization is <10% of total FLOPs. Microbenchmark speedups systematically overstate end-to-end gains—a finding with implications for how systems papers should report performance.

---

Architecture

Rustral is organized as a workspace of 24 focused crates:

Backends (ndarray CPU, Candle CPU/CUDA/Metal, WGPU)
    ↓
Core (Backend trait, TensorOps, ForwardCtx, Module, Parameter)
    ↓
Differentiation (reverse-mode Tape, fused losses)
    ↓
Modeling (Linear, Conv2d, LSTM, Transformer, LLaMA, MoE, GPT-2)
    ↓
Training (SGD/Adam/AdamW, DataLoader, TapeTrainer)
    ↓
Infrastructure (SafeTensors I/O, HF Hub, GGUF, Distributed, Metrics)
    ↓
Deployment (HTTP inference server, model zoo, ONNX export)

The critical design invariant: model code depends only on the Backend and TensorOps traits, not on any concrete backend. Swap the backend constructor and the same model runs on CPU, CUDA, Metal, or WGPU.

---

The ForwardCtx Abstraction

The central design choice. Every forward pass receives:

pub struct ForwardCtx<'a, B: Backend> {
    backend: &'a B,
    mode: Mode,           // Train or Inference
    run_id: u64,          // stable identifier
    shape_policy: ShapePolicy,
    profiler: Option<OperationProfiler>,
}

This eliminates an entire class of bugs. In PyTorch, calling model.train() sets a global flag; forgetting to toggle it produces silent correctness errors. In Rustral, the mode is carried through the call graph, and every operation that behaves differently in training vs. inference receives the mode explicitly. You cannot accidentally run inference with training-mode dropout.

---

Quick Start (5 Minutes)

# Clone and build
git clone https://github.com/skumyol/rustral.git && cd rustral
cargo build --workspace

# Run tests (formats, clippy, full suite — 700+ tests, 0 failures)
./run_tests.sh

# First example: XOR classification (proves the network learns)
cargo run -p rustral-runtime --features training --example tape_xor_classification

# SST-2 sentiment classifier (quick smoke test)
cargo run --release -p rustral-runtime --features training --example sst2_classifier -- --quick

# WikiText-2 language model
cargo run --release -p rustral-runtime --features training --example wikitext2_lm -- --quick

# Benchmark harness (CPU)
python3 scripts/bench/run_all.py --suite rustral --suite candle --repeats 3

GPU Training

# CUDA (requires nvcc + CUDA toolkit ≥ 12.0)
cargo run --release -p rustral-runtime --features "training,cuda" \
  --example sst2_classifier -- --paper --d-model 256 --num-layers 4

# Force Candle CPU backend when CUDA JIT overhead dominates
RUSTRAL_FORCE_CPU=1 cargo run --release -p rustral-runtime --features "training,cuda" \
  --example sst2_classifier -- --paper

LLM Inference

# GPT-2 generation with metrics JSON
cargo run -p rustral-llm --features hf-tokenizers -- \
  generate --model gpt2 --prompt "The capital of France is" --max-new-tokens 32

HTTP Inference Server

cargo run -p rustral-runtime --features training --example save_linear_artifact -- tiny_linear.safetensors
cargo run -p rustral-inference-server -- \
  --artifact tiny_linear.safetensors --bind 127.0.0.1:8080 --in-features 1 --out-features 1 --bias
curl -s -X POST http://127.0.0.1:8080/v1/infer \
  -H 'content-type: application/json' -d '{"input":[[0.25]]}'

---

Backends

// Same model code, three backends
use rustral_ndarray_backend::CpuBackend;
use rustral_candle_backend::CandleBackend;
use rustral_wgpu_backend::WgpuBackend;

let reference  = run(&CpuBackend::default())?;
let optimized  = run(&CandleBackend::cpu())?;
// let on_gpu  = run(&CandleBackend::cuda(0)?)?;

---

Design Philosophy

Rustral favors explicit structure over framework magic.

Framework Comparison

---

EMNLP 2026 Submission

Rustral is the subject of a systems paper submitted to EMNLP 2026 (ARR May cycle). The paper:

• Frames Rustral as a research instrument that makes system internals transparent and reproducible
• Presents the first 6-framework cross-framework operator benchmark under a unified schema
• Proposes Cross-Backend Consistency Regularization (CBCR) — training-time regularization against a frozen reference model
• Demonstrates operator-to-task transfer analysis showing that microbenchmark speedups systematically overstate end-to-end gains

---

Development Status

• Tests: 700+ passed, 0 failures across all crates (excluding experimental WGPU backend from CI)
• WGPU backend: 23/23 correctness tests pass (post binding-index de-confliction)
• LLaMA golden parity: Forward pass logits match HuggingFace Transformers within 2×10⁻³ — CI-automated
• cargo fmt: Clean
• cargo clippy -D warnings: Clean
• Benchmarks: Schema-v2 JSON, CI-validated, 6-framework CPU + 2-framework GPU

Run locally:

./run_tests.sh
# or stricter:
cargo fmt --all -- --check && cargo clippy --workspace --all-targets -- -D warnings \
  && cargo test --workspace --exclude rustral-wgpu-backend

---

Building a Model

use rustral_core::{Backend, ForwardCtx, Mode, Module};
use rustral_nn::{Linear, LinearBuilder, Conv2d, Conv2dConfig};

struct MyModel<B: Backend> {
    conv1: Conv2d<B>,
    fc1: Linear<B>,
    fc2: Linear<B>,
}

impl<B: Backend> MyModel<B> {
    fn new(backend: &B) -> Self {
        Self {
            conv1: Conv2dConfig::new(32, 3, 3).build(backend),
            fc1: LinearBuilder::new(32 * 7 * 7, 128).build(backend),
            fc2: LinearBuilder::new(128, 10).build(backend),
        }
    }

    fn forward(&self, image: &B::Tensor, ctx: &mut ForwardCtx<B>) -> B::Tensor {
        let ops = ctx.backend().ops();
        let x = self.conv1.forward(image.clone(), ctx).unwrap();
        let x = ops.relu(&x).unwrap();
        // ... reshape, fc1, relu, fc2 ...
        x
    }
}

Backend Swap Without Rewriting

fn run<B: Backend>(backend: &B) -> rustral_core::Result<B::Tensor> {
    let model = TinyMlp::new(backend)?;
    let input = backend.tensor_from_vec(vec![0.0, 1.0], &[1, 2])?;
    let mut ctx = ForwardCtx::new(backend, Mode::Inference);
    model.forward(input, &mut ctx)
}

let reference = run(&CpuBackend::default())?;
let fast_cpu  = run(&CandleBackend::cpu())?;

Mixed Precision Training

use rustral_optim::{MixedPrecisionOptimizer, DType};

let optimizer = MixedPrecisionOptimizer::new(Adam::new(0.001))
    .with_dtype(DType::Float16)
    .with_loss_scale(1024.0);

Candle Backend (Recommended for Production)

The candle backend uses the candle-core library for highly optimized CPU execution, plus optional CUDA and Metal paths:

use rustral_candle_backend::CandleBackend;

let backend = CandleBackend::cpu();
// Or with CUDA (requires nvcc / CUDA toolkit):
// let backend = CandleBackend::cuda(0)?;
// Or with Metal on macOS:
// let backend = CandleBackend::metal(0)?;

On large linear layers, candle can be up to ~20x faster than the ndarray backend on CPU.

The benchmark crate also ships CUDA and Metal workload binaries:

cargo run --release -p rustral-bench --features cuda --bin rustral_workloads_cuda -- --repeats 5 --warmup 2
cargo run --release -p rustral-bench --features metal --bin rustral_workloads_metal -- --repeats 5 --warmup 2
# Linux + NVIDIA: CUDA backend tests + runtime (`training,cuda`) + bench JSON integration (`RUSTRAL_TEST_GPU=1`)
./scripts/run_gpu_tests.sh

Flash Attention

use rustral_nn::{FlashAttention, SelfAttentionConfig};

let config = SelfAttentionConfig::new(768, 12);
let flash_attn = FlashAttention::new(&backend, config, 42)?;

Mixture of Experts (MoE)

use rustral_nn::{ExpertLayer, MoEConfig};

let config = MoEConfig::new(512, 64, 2048, 2);
let moe = ExpertLayer::new(&backend, config, 42)?;

---

Examples Gallery

Training and NLP (rustral-runtime)

Layers and Models (rustral-nn)

Candle Backend

---

Documentation

---

Platform Notes

---

Contributing

Good first issues: add an example, improve documentation, add tests for edge cases, fix compiler warnings. See CONTRIBUTING.md.

---

License

MIT OR Apache-2.0 — you may use either.

---

Rustral draws ideas from PyTorch's eager execution, JAX's functional purity, and Rust's ownership model. It aims to stay transparent and hackable alongside crates like Candle and Burn.