eegdino-rs

Rust inference crate for the EEG-DINO foundation model, built on RLX 0.2.5.

EEG-DINO learns robust EEG representations via hierarchical self-distillation on 9 000+ hours of EEG data. This crate provides a numerically verified port of the encoder with NRMSE < 1e-6 against the original PyTorch implementation on RLX CPU, Metal, MLX, and CUDA.

Requirements

• Rust 1.87+
• RLX 0.2.5 (rlx, rlx-cpu, plus optional backend features)
• Weights: weights/eeg_dino_small.safetensors (see Weight conversion)

# Cargo.toml (crates.io)
eegdino = "0.1"
rlx = { version = "0.2.5", default-features = false, features = ["cpu", "cuda"] }

For local development, this repo uses a path dependency to a sibling ../rlx checkout.

Quick start

use eegdino_rs::prelude::*;

let device = parse_device("metal")?; // cpu | metal | mps | mlx | gpu | cuda | rocm
let (mut encoder, load_ms) = EegDinoEncoder::load(
    "weights/eeg_dino_small.safetensors".as_ref(),
    None,
    device,
)?;

let signal = vec![0.0f32; 19 * 2000];
let result = encoder.encode_raw(&signal, 1, 19, 2000)?;
println!("{:?}", result.shape); // [1, 191, 200]

cargo run --release --example infer -- \
    --weights weights/eeg_dino_small.safetensors --device metal

Backends (RLX)

Enable macOS GPU backends:

cargo build --release --features all-backends

NVIDIA / AMD:

cargo build --release --no-default-features --features rlx,rlx-cpu,rlx-cuda
cargo build --release --no-default-features --features rlx,rlx-cpu,rlx-rocm

On Apple Silicon, default features include CPU + Accelerate. Add Metal and MLX with all-backends.

CUDA / ROCm notes (RLX 0.2.5): Attention uses tiled flash kernels for BSHD [B,S,H,D] (EEG-DINO layout). CUDA supports run_slots for zero-copy encode output from the host arena. wgpu matmul parity vs CPU is still being fixed upstream — use Metal or MLX on macOS for production GPU inference today.

Parity

BSHD attention unit tests live in the RLX repo (rlx-ir, rlx-cuda); check_rlx_backends.sh runs them when ../rlx is present.

CUDA production presets

# Latency-oriented
RLX_CUDA_EXEC_MODE=graph cargo run --release --features rlx,rlx-cpu,rlx-cuda \
  --example bench -- --device cuda --batch 64 --only small

# Throughput-oriented (JSON lines for CI)
RLX_CUDA_EXEC_MODE=stream cargo run --release --features rlx,rlx-cpu,rlx-cuda \
  --example bench -- --device cuda --batch 256,512 --json --only small

# Stage breakdown (4 prefix compiles + full encode split)
cargo run --release --features rlx,rlx-cpu,rlx-cuda --example profile_encoder -- \
  --device cuda --batch 16 --stages early

profile_encoder --stages early reports where time goes inside the encoder (transformer dominates on CUDA). Use --stages all for the full pipeline including patch embedding.

Benchmarks (CUDA throughput vs batch size)

# Single steady batch (fastest): CUDA graph replay
RLX_CUDA_EXEC_MODE=graph cargo run --release --no-default-features --features rlx,rlx-cpu,rlx-cuda \
  --example bench -- --device cuda --batch 16 --only small

# Multi-batch sweep: stream exec (graph mode retains VRAM per captured shape)
RLX_CUDA_EXEC_MODE=stream cargo run --release --no-default-features --features rlx,rlx-cpu,rlx-cuda \
  --example bench -- --device cuda --batch 1,2,4,8,16,32,64,128 --only small

Observed on an RTX 4090 (EEG-DINO Small, 19×2000, RLX_CUDA_EXEC_MODE=graph, warmup=5, iters=30):

Large batches with --isolate and RLX_CUDA_EXEC_MODE=stream (one compiled graph per process):

Takeaway: graph mode peaks around batch=64 for steady single-shape serving; stream mode with large batches reaches ~2.1k samples/s on this GPU.

VRAM / OOM when sweeping batch sizes

The encoder caches one compiled RLX graph per (batch, channels, patches) shape. On CUDA/wgpu/ROCm the default is max_cached_shapes = 1: switching batch size drops the previous graph so a sweep like 256,512,1024 does not exhaust VRAM.

let mut enc = EegDinoEncoder::builder()
    .weights(path)
    .device(rlx::Device::Cuda)
    .max_cached_shapes(1) // default on CUDA; use a higher value if you need 2+ shapes hot
    .build()?;
enc.clear_cache(); // optional: force release before a new shape
enc.prewarm_batch_sizes(&[256, 512], 19, 2000)?; // compile + warm graphs ahead of serving

Use encode_batch for batched inference without per-call flatten allocations.

The bench example calls clear_cache() per batch size, auto-enables --isolate when any batch in a multi-size sweep is > 128, and supports --json for CI trends.

cargo run --release --features rlx,rlx-cpu,rlx-cuda --example bench -- \
  --device cuda --batch 256,512,1024 --isolate --only small

To keep several shapes compiled (e.g. production with only B=1 and B=16), set max_cached_shapes(2) or higher.

Model variants

Burn reference (optional)

The original Burn-based implementation remains available for comparison:

cargo build --features burn,ndarray

Types are exported as BurnEegDinoEncoder, etc., when the burn feature is enabled.

Weight conversion

pip install torch safetensors
python scripts/convert_weights.py --checkpoint path/to/model.pt --output weights/

See ABLATION.md for performance notes.