meganeura

Portable neural-network training and inference in Rust. Look, ma, no CUDA!

Warning: project is actively developed. Mostly optimization work, expanding the ops coverage, but with occasional fixes in correctness.

Example

Define a graph, call build_session, train. Meganeura handles autodiff, graph rewrites, WGSL specialization, Naga parsing and validation, and GPU dispatch automatically. The rewrite engine supports a fast deterministic greedy mode and experimental equality-saturation modes.

use meganeura::{Graph, Trainer, TrainConfig, build_session};

let mut g = Graph::new();
let x = g.input("x", &[32, 784]);
let labels = g.input("labels", &[32, 10]);

let w1 = g.parameter("w1", &[784, 128]);
let h = g.relu(g.matmul(x, w1));
let w2 = g.parameter("w2", &[128, 10]);
let logits = g.matmul(h, w2);

let loss = g.cross_entropy_loss(logits, labels);
g.set_outputs(vec![loss]);

// autodiff + graph rewrite + compile + GPU init
let session = build_session(&g);
let mut trainer = Trainer::new(session, TrainConfig::default());
trainer.train(&mut data, /* epochs = */ 10); // data loader: see examples/mnist.rs

A two-layer MLP, trained end to end on the GPU, in one screen.

For local iteration see testing and coverage and debugging the stack.

Why Meganeura

Portable. GPU access is provided by blade-graphics: Vulkan on Linux, Windows, and Android, and Metal on Apple platforms. Mesa's Lavapipe provides a software Vulkan target for headless CI. Doesn't need any runtime.

Lean. Around 50K LOC of Rust+WGSL code in this repository. Automatic shader composition based on kernel archetypes: pointwise, reduction, matmul, convolution, attention. May produce a single 12Mb self-contained binary for deployment.

Fast. Meganeura is pretty fast. It tries to be competitive with vendor-native ML stacks but lands at around 0.5x of their performance today. Your results may very by model, device, precision policy, and the driver. See Inferena tables to get an idea.

How it compares

Meganeura's strong sides are uniform graph, autodiff, compiler, and runtime stack for both training and inference across desktop and edge-class Vulkan/Metal devices.

Install

cargo add meganeura

See the changelog for API and feature changes from 0.2. CI assembles the packaged crate; full registry verification waits on a Blade release that includes this timing API.

Features:

• "hf-hub" to enable HuggingFace downloads
• "models" for built-in models: SmolLM2, SmolVLA, SD_Unet, ResNet, Whisper
• "gguf" for GGUF format loading of weights and graphs

Worked examples live in examples/:

• mnist.rs — MNIST training end to end.
• train_deploy.rs — optimizer-backed training, checkpoint save, and reload into a fresh inference session.
• smollm2.rs — LLM inference with HuggingFace weights.

Current checkpoints store logical tensors without device padding and preflight the restore before mutation. Adam/LaProp moments are allocated only when requested; SGD and forward/backward-only sessions avoid that unused storage. See the checkpoint implementation for format compatibility and restore checks. Resident buffer counts do not measure driver peak memory.

Pretrained models can be loaded from ONNX or NNEF via meganeura::load_onnx(...) / meganeura::load_nnef(...). Both lower through Meganeura’s IR, so the same graph rewrites apply to imported graphs and hand-built ones.

A GGUF file needs nothing alongside it. It carries no graph, but it carries a description — an architecture name and a set of dimensions — and load::gguf reads that into a graph, fills it from the file's own tensors, and uses the tokenizer the file embeds:

use meganeura::load::gguf::{load_gguf, GenerationOptions};

let model = load_gguf(std::path::Path::new("model.gguf"))?;
let mut generator = model.generator(2048)?;
println!("{}", generator.generate("The meaning of life is", &GenerationOptions::default())?);

System requirements

Meganeura runs best when the selected driver exposes hardware-accelerated cooperative matrix operations:

• Vulkan — VK_KHR_cooperative_matrix.
• Metal — simdgroup matrix support.

Falls back to scalar matmul on older hardware. Headless Lavapipe works for CI.

When several adapters are present, select one with its backend-reported numeric device ID (on Vulkan this is normally the PCI device ID, not an adapter ordinal):

MEGANEURA_DEVICE_ID=0x744c cargo run --release --example mnist

Environment variables

All of the environment variables are resolved in SessionConfig::from_env() and never visible to the core modules directly.

Session::tune_with(TuneOptions) enables auto-tuning across multiple dimensions. See the tuning API and whole-step experiment.

Debugging

See testing and coverage for the failure-investigation workflow and performance profiling for timing and capture tools.

Three levels, cheapest first:

• Provenance everywhere. Name values while building (let h = g.matmul(x, w); let h = g.named(h, "blk3.qkv"); — nn layers name their outputs automatically) and the name follows the value through autodiff, rewrites, and fusion into dispatch labels, profiler rows, MEGANEURA_DUMP_PLAN, and NaN reports. Shape panics report the model-builder line that created the bad node.
• Debug sessions. build(&g, SessionConfig::debug()) disables buffer aliasing and keeps everything host-visible: session.read_node_by_name("blk3.qkv") reads materialized values after step() and wait(). Graph rewrites and precision policy are separate controls. session.step_debug() scans primary output prefixes after execution; its first reported NaN/Inf is not a complete root-cause guarantee. Fused-away or aliased values return structured errors.
• Eager evaluation. meganeura::eager::Eager runs the graph you are still building, one eval(&g, node) at a time, on the same kernels the compiled path uses — the PyTorch-style inspect-as-you-go loop. The same graph then compiles unchanged via build_session for training speed.

Profiling

For a repeatable per-dispatch JSON profile:

MEGANEURA_GPU_TIMING=1 \
  cargo run --release --example profile_session -- gap-profile.json

The report retains raw hardware-timestamp samples, selected pipeline variants, workgroup geometry, forward/backward and kernel-family aggregates, device and memory metadata, and the instrumentation overhead relative to normal grouped execution. See structured performance profiling for the Inferena harness and interpretation rules.

For a CPU/GPU timeline:

MEGANEURA_TRACE=trace.pftrace cargo run --example mnist --features profiler

Open the trace in Perfetto:

Citation

Machine-readable author and project metadata is available in CITATION.cff. The paper citation and archival identifier will be added after publication.