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. Graph rewriting uses bounded equality saturation through egglog.

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 class labels, Graph::cross_entropy_loss_indices takes an input_u32 of shape [batch]. Upload indices with Session::set_input_u32; the loss and gradient read these indices directly. Use cross_entropy_loss for dense or weighted targets.

For local iteration see testing and coverage and debugging the stack. The shader generation investigation compares WGSL, synaga and other shader construction approaches.

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

	GPU backends	Training	Approach
Meganeura	blade-graphics (Vulkan, Metal)	yes	graph IR + rewrites + specialized WGSL
Candle	CUDA, Metal, CPU	yes	tensor API, native kernels
Burn	CubeCL: CUDA, ROCm, Metal, Vulkan, WebGPU; CPU paths	yes	modular backends, JIT fusion
tch-rs	CUDA, CPU (via libtorch)	yes	PyTorch FFI bindings

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.

Contents

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.
• gemma4.rs — Gemma 4 GGUF decode vs llama.cpp.

The matched GGUF diagnostic compares Vulkan latency and logits against a pinned llama.cpp checkout, including host readback.

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())?);

load_gguf memory-maps files opened from disk, so model weights fault in on demand instead of first being copied into an equally large heap allocation. The in-memory load_gguf_bytes and load_gguf_shared entry points retain their owned-buffer behavior.

CompileOptions::quantized_activations quantizes the GEMV activation row to Q8_1 and runs the inner product on integer dot products, following llama.cpp's vec_dot_*_q8_1 kernels. It is on by default: a model that ships quantized weights is decoded quantized, so its activations join them, and every format with a kernel layout for it — GGML Q4_0, Meganeura Q8, and the K-quants Q4_K, Q5_K, Q6_K and Q3_K — gets it automatically. It is the only kernel switch here that changes results rather than the route to them, so it is never selected by measurement; set it false to keep f32 activations. Where the device reports shader_integer_dot_product (Metal, and Vulkan 1.3 devices exposing VK_KHR_shader_integer_dot_product) the dot products are the hardware dot4I8Packed (DP4A); otherwise an exact scalar expansion of the same integer arithmetic runs. Q2_K and Q8_K are listed in the inventory but not read, and quantized embedding tables have no gather variant.

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.

For joint graph and implementation search with representative inputs, use train::build_measured with initialization and numerical-qualification callbacks. It replaces live attention/submission retuning; ordinary build does not execute an uninitialized model. See compiler search.

Variable	Effect
MEGANEURA_DISABLE_COOP	Force the portable scalar matmul path (regression diagnosis).
MEGANEURA_COOP_F16	Allow f16-input cooperative tiles even for precision-sensitive derivative work; requires numerical validation. Default Auto protects derivatives and can already use f16 tiles for forward work.
MEGANEURA_FLASH_FWD_COOP=0	Disable only cooperative flash-attention forward.
MEGANEURA_FLASH_BWD_COOP	Enable the experimental reduced-precision flash backward.
MEGANEURA_NO_ALIAS	Disable buffer lifetime aliasing (every value gets its own allocation).
MEGANEURA_NO_DEVICE_LOCAL	Keep all buffers host-visible.
MEGANEURA_POISON	Fill non-parameter buffers with NaN at session build, so reads of memory nothing wrote show up.
MEGANEURA_SERIAL_DISPATCH	One compute pass per dispatch — serial execution for bisection.
MEGANEURA_NO_WINOGRAD	Skip the Winograd rewrite of 3×3 stride-1 convolutions and their input gradients; its selection heuristic weighs channel counts only, so this measures which side of it a workload belongs on.
MEGANEURA_PIN_BUFS=3,25-40	Force-pin logical buffers to bisect aliasing corruption.
MEGANEURA_ARENA_CHUNK_BYTES	Largest parameter-arena chunk in bytes; smaller values spread the optimizer over more dispatches.
MEGANEURA_DUMP_PLAN	Dump dispatch order, provenance, and the alias map at build.
MEGANEURA_DUMP_WGSL=<dir>	Write every generated shader into <dir>.
MEGANEURA_OPTIMIZER	Rewrite mode: off | egglog-windowed | egglog-outlined (default) | egglog-whole. Legacy greedy maps to outlined egglog.
MEGANEURA_EGRAPH_COST	Extraction objective: ast-size | tensor-traffic.
MEGANEURA_EGRAPH_CUTOFF=<n>	Saturation segment-size ceiling (default 300).
MEGANEURA_GREEDY_PACK_SWIGLU=0	Skip packing consecutive SwiGLU ops into one parameter buffer during graph optimization (legacy variable name).
MEGANEURA_DEVICE_PARAMETERS	Experimental placement of unaliased parameter buffers on the device: 1 → device-transient, device-buddy → device. Default is host-visible.
MEGANEURA_REUSE_UPLOAD	Reuse one staging buffer across set_parameter uploads instead of restaging per parameter.
MEGANEURA_TUNE	Opt-in bounded matmul, convolution and GEMV search at build (SessionConfig { tune: true }), using private scratch. Scalar tiles and GEMV shapes also support reduced-storage weights.
MEGANEURA_FLASH_EPT_CAP=<n>	Flash forward elements-per-thread cap (power of two ≥ 2).
MEGANEURA_FLASH_THREADS=<n>	Flash forward workgroup width: 128 or 256.
MEGANEURA_FLASH_KEYS=<n>	Flash forward shared key/value tile: power of two up to 16, capped by device shared memory.
MEGANEURA_FLASH_INTERLEAVE	Interleave head dimensions across forward attention lanes.
MEGANEURA_FLASH_GRAD_Q_EPT_CAP=<n>	EPT cap for flash dQ backward.
MEGANEURA_FLASH_GRAD_KV_EPT_CAP=<n>	EPT cap for fused flash dK/dV backward.
MEGANEURA_FLASH_BWD_EPT_CAP=<n>	Shared fallback cap for both flash backward kernels.
MEGANEURA_DEVICE_ID=0x744c	Adapter selection by numeric device id.
MEGANEURA_GPU_TIMING	Enable hardware timestamp pools (set before context creation).
MEGANEURA_GPU_CAPTURE	Enable Blade's native-tool labels and shader debug information before context creation; independent of GPU timing.

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.