Skip to content

kvark/meganeura

Folders and files

NameName
Last commit message
Last commit date

Latest commit

 

History

316 Commits
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 
 

Repository files navigation

meganeura

CI Docs Crates.io

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.

logo

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.rs

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

Why Meganeura

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.

How it compares

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.

Install

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.

System requirements

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

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

Decimal 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.

Profiling

MEGANEURA_TRACE=trace.pftrace cargo run --example mnist

Open the trace in Perfetto:

perfetto trace

Contributing

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.

About

Blade-powered NN training and inference

Topics

Resources

Stars

Watchers

Forks

Releases

Packages

Contributors

Languages