Portable neural-network training and inference in Rust. Meganeura uses Vulkan on Linux, Windows, and Android and Metal on Apple platforms. It does not require CUDA, ROCm, or Python at runtime.
Status: actively developed; APIs and benchmark methodology are still in motion. Current workloads include SmolLM2, a SmolVLA action expert, ResNet-50, the Whisper-tiny encoder, and a scaled, timestep- and text-conditioned latent-diffusion U-Net. Issues and pull requests are welcome.
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.rsA two-layer MLP, trained end to end on the GPU, in one screen.
Fast. Meganeura is competitive with vendor-native ML stacks on selected workloads while retaining one Vulkan/Metal implementation. Results vary substantially by model, device, precision policy, and driver. The historical Inferena tables are useful exploratory data, but they predate the audited paper protocol and should not be treated as publication-grade comparisons. The new protocol reports raw samples, uses matched workloads and full forward-plus-backward timing, and separates strict f32 from reduced-input accelerated modes.
PyTorch CUDA currently leads several RTX 5080 workloads, especially training. Meganeura's strongest result is therefore not universal speed superiority; it is how much of that performance can be reached through a portable execution stack that also runs on AMD, Intel, and Apple GPUs.
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. The compute stack does not require CUDA or ROCm, although performance and feature availability still depend on each vendor's driver.
Composable. A small set of
kernel archetypes
— pointwise, reduction, matmul, convolution, and attention — compose into
specialized GPU shaders at compile time. The rewrite set recognizes
equivalent fused forms (for example, x * sigmoid(x) → SiLU and
SiLU(gate) * up → SwiGLU).
It can run either as a deterministic greedy pass or through equality
saturation with a traffic-aware extraction cost. Current ablations find the
same selected graph for the benchmark rewrite set, so equality saturation is
research infrastructure rather than a claimed source of runtime speedup.
Consolidating the remaining hand-written WGSL variants into parameterized
generators is active work.
| GPU backends | Training | Approach | |
|---|---|---|---|
| Meganeura | blade-graphics (Vulkan, Metal) | yes | graph IR + rewrites + specialized WGSL |
| Candle | CUDA, Metal, CPU | limited | eager tensors, hand-written kernels |
| Burn | CUDA, wgpu, NDArray, LibTorch | yes | modular multi-backend |
| tch-rs | CUDA, CPU (via libtorch) | yes | PyTorch FFI bindings |
Meganeura's wedge is a uniform graph, autodiff, compiler, and runtime stack for both training and inference across desktop and edge-class Vulkan/Metal devices.
cargo add meganeura
Worked examples live in examples/:
mnist.rs— MNIST training end to end.smollm2.rs— LLM inference with HuggingFace weights.
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.
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 mnistDecimal IDs are accepted too. MEGANEURA_DISABLE_COOP=1 forces the portable
scalar path for regression diagnosis; MEGANEURA_FLASH_FWD_COOP=0 and
MEGANEURA_FLASH_BWD_COOP=0 disable only cooperative flash attention.
Device-local intermediate storage and lifetime aliasing are default-on, with
MEGANEURA_NO_DEVICE_LOCAL=1 and MEGANEURA_NO_ALIAS=1 as diagnostic escape
hatches.
Hosted CI executes Linux Vulkan (Lavapipe) and macOS Metal tests and compile-checks Windows. Generated Vulkan shaders, including the f16 cooperative-matrix path used on supported NVIDIA and AMD adapters, are also validated and translated to SPIR-V offline. Real AMD and NVIDIA hardware testing remains part of release qualification because software drivers cannot reproduce vendor-driver behavior.
MEGANEURA_TRACE=trace.pftrace cargo run --example mnist
Open the trace in Perfetto:
Early project, small API surface, small community — a good time to show up. Open an issue before starting anything sizeable so we can align on direction.