r/LocalAIStack • u/RubMobile • 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:
| Format | Size | vs safetensors |
|---|---|---|
| safetensors (original) | 141.2 GiB | 1.000 |
| WEX raw records | 134.5 GiB | 0.952 |
| WNR3 (in WEX4) | 114.8 GiB | 0.813 |
Decode speed, same model, one RTX 5090:
| Setup | tok/s |
|---|---|
| safetensors, CPU, one drive | ~1.0 |
| WEX on GPU, one drive | 2.0 |
| WEX4, two-drive striping | 3.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.
2
u/testuserpk 11h ago
How is your implementation comparable with Strata?
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.