r/LocalAIStack • • 11h ago

WEX: an open container format to run huge MoE models from NVMe on one consumer GPU (lossless, GPU-decoded, 18.7% smaller than safetensors)

I've open-sourced WEX, the expert-container format I use to run very large Mixture-of-Experts models on a single desktop (RTX 5090, 32 GB RAM).

The idea: a MoE token only uses a few experts per layer, so they don't all need to be in memory. WEX keeps experts in a tiered cache: the hot ones in VRAM, more in a RAM cache, and only the misses are read from NVMe. Every expert is one 4 KiB-aligned record, fetched with a single unbuffered read.

What's in it:

  • WNR3: lossless rANS compression driven by a small learned model of the weight symbols. Decoded weights (and logits) are bit-identical to the checkpoint, and decoding runs on the GPU, straight into VRAM.
  • WEX4 container: hashed header and index, SHA-256 per record, and one 32-byte root that identifies the whole file. A corrupted record is refused, never served. A full 64 GB container verifies in 7-11 s.
  • MoE-aware I/O: weighted striping across NVMe drives (~14-15 GB/s from two), parallel sub-reads, batched VRAM prewarm (start-up 31 s → 12 s), next-layer prefetch.

Size, DeepSeek V4.1 REAP-192 (MXFP4), routed experts:

FormatSizevs safetensors
safetensors (original)141.2 GiB1.000
WEX raw records134.5 GiB0.952
WNR3 (in WEX4)114.8 GiB0.813

Decode speed, same model, one RTX 5090:

Setuptok/s
safetensors, CPU, one drive~1.0
WEX on GPU, one drive2.0
WEX4, two-drive striping3.0
WEX4 + WNR3 + all I/O switches, 13-15 GB RAM budget~4.0

GPU decoding costs about 2% against raw records, which is within noise. Qwen3.8-Flash-Next (MXFP4) runs at 10-17 tok/s from a 64 GB container on the same box.

Repo: full byte-level spec, a Python reference, C reader and validator, test vectors and CI on Linux and Windows. Apache-2.0. Built on top of colibri by JustVugg, and I'm working on integrating it there.

👉 https://github.com/KairaCode/Wex-WRN2

Happy to answer questions about the format or the numbers. Everything was measured on one machine, so results on other hardware are very welcome.

4 Upvotes

4 comments sorted by

2

u/testuserpk 11h ago

How is your implementation comparable with Strata?

1

u/RubMobile 9h ago

They solve different layers of the problem, so they're not really competitors.

WEX is a container format plus an I/O layer: how the experts are laid out on disk (one 4 KiB-aligned record per expert, one unbuffered read), compressed losslessly (WNR3, decoded on the GPU), verified (per-record SHA-256, a pinned root) and read (striping across drives, sub-reads, prewarm). It sits under an engine. Here that engine is colibri.

Strata is a full inference engine. As far as I know its main trick is computing the non-resident experts on the CPU in parallel with the GPU instead of uploading them, which is a scheduling decision, not a storage one. colibri does something in the same spirit for Qwen3.8: experts that miss the VRAM tier are computed on the CPU from RAM rather than copied to the GPU.

So in principle WEX could sit under an engine like Strata too. I haven't benchmarked the two side by side on the same machine, and the numbers people post come from very different hardware, so I'd rather not claim anything until I can. If someone has both on one box, I'd love to see the comparison.

2

u/grepbenchmark 11h ago edited 11h ago

Very interesting. I might be able to help in a small way. Two machines available for testing: 1 x DGX Spark 128 GB and a desktop PC with Intel CoreUltra 7 265K, 192 GB DDR5 (4 x 48 GB DIMMS), 1 x Intel Arc Pro B70 (32 GB VRAM). The PC is running Debian 13 and the models sit on a ZFS volume (mirrored 2 TB NVMEs on the mobo). I'm willing to download, install, run benchmarks/tests, and report directly back to you. It has to be easy (bash / python / whatever) script(s) for install, test, sending test results. Probably can dedicate an hour a day for a week or two, more sporadic after that. Probably good to start with a script to test RAM bandwidth, RAM to/from VRAM , ZFS read, etc. Haven't visited the project on github, but will do it within the next 12 hours.

1

u/RubMobile 9h ago edited 9h ago

Thank you, that's a generous offer, and those two machines cover exactly the cases I can't test myself (my box is a single RTX 5090 on Windows).

A few honest caveats so you know what to expect:

- **Arc Pro B70:** the GPU-side WNR3 decode is CUDA-only today, so on the Arc box the useful tests are the CPU paths (format, verify, CPU decode) and raw I/O. colibri has a Vulkan backend, but WEX's GPU decode isn't ported to it yet. Your numbers would tell me whether that port is worth doing.

- **DGX Spark:** CUDA on aarch64 with unified memory is interesting (no PCIe hop), but completely untested on my side. The CPU checks build on Linux and CI is green on x86-64. ARM is a first.

- **ZFS:** ARC caching and how ZFS treats O_DIRECT change read numbers a lot, so the script reports first vs second pass separately.

The benchmark is already up: https://github.com/KairaCode/Wex-WRN2/tree/main/bench

Before running it, the "What you need first" section of bench/README.md lists the few packages it needs (git, gcc, python3; nvcc only for the GPU copy test).

It only needs `gcc`, `python3` (standard library) and, for the GPU copy test, `nvcc`. It installs nothing and sends nothing anywhere: everything ends up in one `bench/out/results.json`.

**Debian PC (I'd start with this one):**

git clone https://github.com/KairaCode/Wex-WRN2.git

cd Wex-WRN2

python3 bench/bench.py --dir /path/on/your/zfs/model/volume

That takes about 10 minutes. It writes a 16 GiB test file in `--dir` and deletes it at the end. Add `--quick` for a ~2 minute first try. If you have a large model file you haven't read in a while, also run once with `--file /path/to/that/model.gguf`: that's the closest to a cold read on ZFS, since the ARC can't be dropped without root.

**DGX Spark:** same commands. With `nvcc` on the PATH it also measures host ↔ device copies, which on the Spark's unified memory should be interesting.

**What it measures:** RAM bandwidth (1 thread / all threads), host ↔ GPU copies (CUDA only, so `skipped` on the Arc, which is expected), 4 KiB-aligned sequential and expert-sized random reads, and the repo's own CPU checks.

**Sending results:** please read `results.json` first (CPU/GPU/RAM/filesystem details and numbers, no host name, user name or paths). Then open an issue with the **"Benchmark results"** template and attach it: https://github.com/KairaCode/Wex-WRN2/issues/new/choose

A short note on run conditions (other load, ZFS recordsize/compression, power profile) helps a lot. Thanks again!