Streaming Mixture-of-Experts Weight Tensors Directly from NVMe Storage

Mixture-of-Experts (MoE) architectures split feed-forward networks into discrete, specialized sub-networks called experts. A routing gate calculates token-to-expert affinity scores and activates only a small subset of experts per token during inference. While this sparse activation pattern keeps per-token computational cost low, the physical memory footprint remains massive because all expert weights traditionally sit resident in RAM or GPU VRAM.

Problem: Memory Starvation in Sparse Architectures

Dense transformer models execute every weight matrix for every single token. Memory allocation is straightforward because model weights remain hot throughout the forward pass. MoE models break this pattern. An 8x7B architecture contains roughly 47 billion parameters across all layers, yet a single forward pass routes each token through only two active experts, computing roughly 13 billion parameters per step.

Traditional inference runtimes like vLLM, TensorRT-LLM, and standard llama.cpp implementations demand that the entire 47-billion parameter matrix stay mapped inside fast memory. On a server with 32 GB of system RAM, attempting to load an uncompressed float16 MoE model triggers the Linux Out-Of-Memory killer instantly. Even running an 8-bit quantized model strains workstation memory buses. Inactive experts sit completely idle for dozens of consecutive tokens, consuming tens of gigabytes of physical memory without performing floating-point operations.

Operating system swap spaces offer an emergency fallback, but standard kernel page faults destroy throughput. When the router picks an expert whose memory pages reside on disk, the execution thread blocks synchronously while the kernel reads 4 KB pages from storage. Inference rates collapse from 35 tokens per second to less than 0.5 tokens per second.

Existing Solutions: Quantization and Synchronous Offloading

Engineers currently tackle this memory wall through aggressive weight quantization and host-to-device memory staging. Quantizing weights to 4-bit representations using formats like GPTQ or EXL2 shrinks total storage requirements. An 8x7B model drops from 90 GB down to roughly 28 GB. While this compression fits medium-tier GPUs, it fails when scaling to modern sparse architectures like 8x22B or DeepSeek MoE models that exceed 140 GB even at low bit widths.

Naive offloading engines read expert weight tensors from disk or system RAM on demand over PCIe. These implementations use standard POSIX read calls or basic memory-mapped files without prefetching hint flags. The process incurs multiple memory copies: disk blocks travel into kernel page cache, get copied into user-space buffers, and finally move across the PCIe bus into GPU memory. This serialization introduces massive pipeline stalls.

Sequential read bandwidth on PCIe 4.0 and PCIe 5.0 NVMe drives reaches 7 GB/s to 14 GB/s. The bottleneck in sparse model offloading is not raw disk throughput, but rather the latency of synchronous kernel page faults and redundant buffer allocations.

Direct memory mapping with asynchronous kernel advice reduces buffer duplication. By mapping raw safetensors or GGUF files directly into the virtual address space with shared memory flags, the application bypasses intermediate user buffers:

// Map model weights file directly into virtual address space
int fd = open("model-moe-q4.gguf", O_RDONLY | O_DIRECT);
size_t file_size = get_file_size(fd);
void *weights = mmap(NULL, file_size, PROT_READ, MAP_SHARED, fd, 0);

// Notify kernel of sequential access patterns before layer execution
madvise(weights + expert_offset, expert_bytes, MADV_WILLNEED);

// Release physical pages immediately after matrix multiplication completes
madvise(weights + expert_offset, expert_bytes, MADV_DONTNEED);

Future with AI: Predictive Routing and Asynchronous IO Rings

The next evolution in local sparse model execution pairs asynchronous kernel primitives with predictive routing gates. Linux io_uring interfaces enable ring-buffer based, zero-copy storage transfers without context switches. Inference engines submit batch read requests for target expert tensors directly to the NVMe controller queue, writing payload blocks directly into pinned direct memory access (DMA) host buffers.

The critical performance advance comes from lookahead expert prediction. Transformer attention layers evaluate token embeddings several layers before the token arrives at the feed-forward network blocks. By training an auxiliary linear classifier or evaluating router softmax distributions two layers in advance, the runtime predicts required expert indices ahead of execution time.

  • Early routing gate calculates top-k expert IDs at layer N minus 2.
  • An asynchronous io_uring SQE (submission queue entry) requests expert weight tensors from NVMe storage.
  • Weights arrive in host pinned memory while layers N minus 1 compute self-attention.
  • Compute kernels execute GEMM operations on fetched expert weights with zero pipeline delay.

This predictive streaming approach eliminates the need to hold hundreds of gigabytes of inactive weights in RAM. Edge servers and developer workstations can run large sparse architectures at near-native compute speeds directly from fast local storage.

Press Cmd K to search برای جستجوی سایت از Cmd+K استفاده کنید