PaperlayerView exact revision on arXiv

Source-linked AI summary

Every Expert Counts: ExactMoE for Memory-Efficient W4A16 Inference

Amjad Saab

arXiv:2608.15383v1cs.LG

TL;DR

Sparse MoE deployment still requires storing and moving inactive experts, creating GPU-capacity and data-movement bottlenecks. ExactMoE combines complete-expert W4A16 quantization with kernel-native host storage, configurable GPU residency, and fused execution; on OLMoE, cache 16 reduced reserved GPU memory 87.04% while retaining 81.85% of BF16 decode throughput.

  • Problem

    Sparse MoE deployment still incurs the cost of storing and moving inactive experts, creating GPU-capacity, host-to-device traffic, and quantization bottlenecks.

  • Method

    ExactMoE integrates expert-only W4A16 quantization, kernel-native host storage, configurable GPU residency, and fused grouped execution while retaining all experts.

  • Results

    87.04% lower peak reserved GPU memory and 81.85% of BF16 decode throughput were achieved with a 16-slot cache on OLMoE.

  • Takeaways & Limitations

    ExactMoE identifies a practical memory–transfer–throughput frontier for complete-expert MoE inference on the evaluated OLMoE system.

  • Takeaways & Limitations

    Performance depends on routing locality, cache capacity, host-to-GPU bandwidth, and transfer latency, with limited expert reuse potentially eliminating offloading’s throughput benefit.

Abstract

from arXiv · show

Sparse mixture-of-experts (MoE) language models reduce arithmetic by activating only a small subset of experts per token, yet deployment still requires storing and moving the full expert bank. We present ExactMoE, an inference design that applies symmetric group-128 four-bit weight quantization only to routed experts, stores those experts in kernel-native MARLIN form in pinned host memory, and executes all selected experts through a configurable GPU-resident slot cache and fused grouped MoE kernels. The router, attention, embeddings, normalization layers, and language-model head remain in BF16. "Exact" refers to complete expert availability and an unchanged top-k routing procedure: no expert is pruned, substituted, or forced to execute on the CPU. It does not imply numerical identity with the BF16 model. On OLMoE-1B-7B-0924-Instruct, evaluated on a single NVIDIA L4, a 16-slot configuration reduces peak reserved GPU memory from 14.168 to 1.836 GiB (87.04%) while retaining 81.85% of BF16 decode throughput. A fully resident 64-slot configuration reaches 31.923 tokens/s versus 21.662 tokens/s for BF16 while reserving 4.061 GiB. Across 12,450 zero-shot multiple-choice questions, ExactMoE obtains 70.3534% normalized accuracy versus 70.8996% for BF16, retaining 99.23% of the baseline accuracy. In a matched 16-token ablation, fused grouped execution is 1.97x as fast as a sequential W4 reference. These results identify a practical memory-transfer-throughput frontier for complete-expert MoE inference.

1 Introduction

ExactMoE addresses the deployment costs of sparse MoE models by combining routed-expert W4A16 quantization with kernel-native host/GPU storage and fused, cache-aware execution. Its systems contribution is a complete-expert runtime that preserves routing while exposing a measured capacity frontier.

  • Motivation: Sparse MoE deployment still requires storing and moving inactive experts, creating bottlenecks in GPU capacity, host-to-device traffic, or quantization error.Conditional computation reduces arithmetic and training cost but does not remove the deployment cost of the full expert bank.
  • Method: ExactMoE quantizes routed experts once into the packed MARLIN format consumed by the GPU kernel and stores that representation in pinned host memory and a configurable GPU cache.Cache misses copy experts without dequantization or repacking.
  • Method: At each sparse layer, ExactMoE resolves active experts in one batched transfer, partitions over-capacity active sets into resident waves, and accumulates fused multi-expert kernel outputs.An explicit global-to-slot map ensures that all selected experts contribute to the result.
  • Contribution: The contribution integrates expert-only W4A16 quantization, LRU-style caching, MARLIN, grouped fused execution, and a common kernel-native host/GPU representation into a complete-expert runtime.The router and dense model components remain in BF16 while only routed experts are quantized.
  • System Design: ExactMoE separates offline representation building, router-preserving token grouping, and tiered fused execution so method invariants remain distinct from empirical performance claims.This separation structures the conversion and inference paths described in Figure 1.

2 Related Work

Prior work spans sparse MoE architectures, expert quantization, tiered expert caching and offloading, and weight-only kernels and serving runtimes. ExactMoE combines these directions through MARLIN-packed W4A16 experts, complete expert availability, GPU execution, and cache management.

  • Sparse mixture-of-experts: Sparse MoE models differ in checkpoint layout, expert geometry, routing policy, attention architecture, and auxiliary components, complicating portable expert deployment.The surveyed examples include Switch Transformers, OLMoE, and Qwen MoE families.
  • MoE quantization: Expert quantization research includes low-bit expert weights, precision allocation, sensitivity- or frequency-based policies, low-rank compensation, and runtime precision adaptation.MoQE specifically showed that expert-only quantization can reduce model size and latency.
  • Expert offloading and caching: Prior offloading systems use activation locality, virtual experts, CPU/GPU coordination, or CPU/GPU/I/O pipelining to manage expert residency and transfers.ExactMoE shares the tiered-residency objective but uses one W4 representation in every tier and executes every selected expert on the GPU.
  • ExactMoE positioning: ExactMoE contributes expert-cache management and complete-expert grouped execution while combining MARLIN-packed weights with vLLM’s fused MARLIN MoE primitive.The related systems were not reproduced under the same hardware and workload protocol, so ExactMoE does not claim superiority over them.
  • Weight-only kernels and serving runtimes: GPTQ and AWQ established post-training weight quantization, while MARLIN provides mixed-precision W4A16 kernels and vLLM provides high-throughput serving with PagedAttention.ExactMoE uses MARLIN-format packed weights and vLLM’s fused MARLIN MoE primitive.

3 Method

ExactMoE quantizes routed expert matrices into groupwise INT4 MARLIN representations, stores them in pinned host memory, and serves them through GPU slot caches and fused grouped kernels. Its design preserves the router, top-k rule, and complete expert index set while supporting cache misses and multi-wave execution.

  • Routing-faithful expert quantization: ExactMoE preserves the router’s top-k rule and every expert index while replacing only gate, up, and down expert matrices with groupwise INT4 representations.Quantization can change hidden states, so later routing decisions may differ from a BF16 rollout.
  • Groupwise INT4 packing: 128-weight groups use per-group scales and quantized values, which are packed with FP16 scales into MARLIN-native tensors offline.Gate and up projections are concatenated for the fused kernel, while down projections are packed separately; all-zero groups use zero scale and quantized values.
  • Host-resident expert pool: 3.09 GiB is the measured packed host-pool size for OLMoE’s 6.44B routed-expert parameters.The host pool stores final packed kernel representations in pinned memory, so cache misses copy them asynchronously without materializing FP16 experts.
  • GPU slot cache: Each sparse layer uses C expert slots with k ≤ C ≤ E, mapping global expert IDs to resident slots and replacing unprotected least-recently-used entries on misses.CUDA events prevent a slot from being overwritten before its preceding computation finishes; when C = E, all experts are preloaded and replacement is disabled.
  • Wave-based grouped execution: Active experts are processed in waves of at most C, with each wave loading experts, building an ID map, running fused MARLIN kernels, accumulating outputs, and recording completion events.The production path uses one fused gate/up grouped GEMM and one down grouped GEMM per wave.

4 Experimental Setup

The experiments use a controlled OLMoE-1B-7B-0924-Instruct comparison between BF16 and ExactMoE on one NVIDIA L4 GPU. They vary expert-cache capacity, measure runtime and memory under fixed protocols, and evaluate zero-shot likelihood-based accuracy across three benchmarks.

  • Model and comparison: OLMoE-1B-7B-0924-Instruct has 16 sparse layers, 64 experts per layer, top-8 routing, hidden size 2048, and 6.44 billion routed-expert parameters.BF16 and ExactMoE use the same checkpoint, tokenizer, SDPA attention backend, prompt set, and greedy decoding policy.
  • Hardware and software: Experiments ran on one NVIDIA L4 with 22.0 GiB usable GPU memory and 53 GiB host memory, using PyTorch 2.11.0+cu130, CUDA 13.0, Transformers 5.14.1, and vLLM 0.26.0.Packing aligned with a specified Marlin commit, and grouped execution used vLLM’s fused MARLIN MoE primitive.
  • Runtime evaluation: Cache capacities were C ∈{8, 16, 32, 64}, with C = 16 as the primary offloaded comparison and capacity 64 fully resident.Runtime tests used 10 fixed prompts, three repetitions, and 128 generated tokens per request, reporting decode throughput, TTFT, peak reserved GPU memory, cache events, and transferred bytes.
  • Batch evaluation: The batch sweep uses fixed batches {1, 2, 4, 8} and 32 generated tokens per request, with one descriptive measurement per batch size and no request scheduler.It measures ordinary batched generation rather than continuous-batching server behavior.
  • Accuracy evaluation: Evaluation covers 570 ARC-Easy, 1,838 PIQA, and 10,042 HellaSwag validation questions, totaling 12,450 paired questions.Normalized accuracy uses highest mean token log-likelihood, while raw accuracy uses highest summed log-likelihood; teacher-forced likelihood is additionally reported on 120 WikiText-2 passages.
  • Model and comparison: The controlled comparison isolates W4A16 expert representation and cache capacity on one checkpoint and one GPU against the original BF16 checkpoint and an internal sequential W4 reference.Prompt repetitions estimate within-process variability rather than fresh-process or cross-machine variance.

5 Results

ExactMoE establishes a memory–throughput frontier in which larger expert caches reduce transfers and improve decode speed, while full residency reaches a quantized-compute endpoint above BF16 throughput. It preserves nearly all BF16 multiple-choice accuracy, doubles matched sequential-reference speed with identical greedy outputs, and scales batched throughput through grouped expert GEMMs.

  • Capacity frontier: 87.04% lower peak reserved GPU memory at 16 slots accompanies 17.731 tokens/s decode throughput, or 81.85% of BF16.This configuration incurs 21.44 GiB of expert traffic per measured request and 290 ms median TTFT.
  • Capacity frontier: 31.923 tokens/s at full residency is 47.4% above BF16 while reserving 4.061 GiB.Full residency is the quantized compute endpoint of the frontier, not an offloading result.
  • Capacity frontier: 21.64 ± 0.18 tokens/s is the BF16 mean across 30 persistent-warm observations, versus 15.90 ± 0.51, 17.84 ± 0.57, 21.69 ± 0.56, and 31.86 ± 0.28 for capacities 8, 16, 32, and 64.These repeated prompt measurements characterize within-process variability rather than independent machine replications.
  • Quality: 99.23% of BF16 normalized multiple-choice accuracy is retained across the 12,450-example aggregate.The absolute aggregate change is −0.546 percentage points, with a paired 95% bootstrap interval of [−0.964, −0.137] points and exact McNemar p = 0.0112.
  • Kernel ablation: 1.97× speedup results from grouped fused execution versus the retained sequential expert reference at a matched 16-slot cache.The paths reach 18.43 and 9.36 tokens/s respectively, generate exactly the same 16 greedy tokens, and have maximum last-token-logit difference 0.421875.
  • Batched throughput: 90.02 tokens/s at batch 8 versus 15.74 tokens/s at batch 1 demonstrates aggregate W4A16 throughput scaling as routes are coalesced into grouped expert GEMMs.At batch 8, W4A16 exceeds measured BF16 aggregate throughput while using substantially less GPU memory; each batch point is measured once without request scheduling.

6 Discussion

ExactMoE’s operating point trades GPU memory against expert-transfer overhead, enabled by kernel-native MARLIN storage and complete routed-expert execution. Its 99.23% accuracy retention is a strong trade-off, but not numerical identity with BF16.

  • Operating point: C = 16 moves a median 21.44 GiB per 128-token request and runs 18.1% slower than BF16 decode on the L4.Capacity 32 nearly matches BF16 throughput while retaining an 81.8% reserved-memory reduction; full residency suits 4–5 GiB allocations.
  • Kernel-native storage: 142.7 seconds is required for OLMoE W4 loading and packing versus 4.71 seconds for BF16 loading in the measured environment.The host pool stores final MARLIN data, including packed qweights and FP16 scales, avoiding per-miss transformation and dequantization workspace.
  • Meaning of “Exact”: ExactMoE preserves the router and executes every selected expert, unlike pruning, substitution, or CPU-tail schemes that change expert availability or execution semantics.Quantized experts can alter hidden states, later router decisions, and BF16 logits; “exact” denotes a systems property rather than numerical identity.
  • Accuracy trade-off: 99.23% of baseline accuracy is retained, but the paired aggregate test detects a small reduction.The paper reports retention and absolute change together and recommends preregistering task weighting and a non-inferiority margin in future studies.

7 Limitations

ExactMoE’s benefits depend on host-memory capacity, transfer behavior, routed-parameter share, and expert reuse. Its numerical, hardware, packing, update, and multi-GPU constraints limit portability and generality.

  • System and performance constraints: Performance depends on routing locality, cache capacity, host-to-GPU bandwidth, and transfer latency; limited expert reuse can eliminate offloading’s throughput benefit.Selected experts absent from the GPU must be transferred from the packed host pool.
  • Memory constraints: Memory savings depend on the routed-expert parameter fraction, while the complete packed expert pool must fit in host memory.Attention, routing, embeddings, normalization, shared components, and the language-model head remain in dense precision.
  • Numerical and quality limitations: W4A16 does not preserve numerical identity with BF16, and quantization may change expert outputs, probabilities, hidden states, and later routing decisions.Quality retention may vary across models, tasks, and experts despite unchanged router parameters and top-k procedure.
  • Hardware and portability constraints: The implementation requires MARLIN-compatible dimensions, grouping constraints, and NVIDIA GPUs, while new architectures require checkpoint-tensor mapping and potentially different packing or fused kernels.Other expert layouts may require padding.
  • Deployment scope: ExactMoE targets fixed-weight inference on one device; expert updates require repacking, and multi-GPU placement, distributed cache coordination, and continuous-serving schedulers need additional support.These capabilities are outside the current implementation target.

Reproducibility Statement

The reproducibility record fixes the experimental configuration and preserves detailed artifacts linking reported tables and figures to measurements. The archive includes data, predictions, statistics, logs, runtime files, and checksums.

  • The experimental record fixes model and dataset revisions, random seed, software stack, MARLIN implementation, attention backend, cache capacities, prompt manifests, generation settings, and statistical procedures.
  • The archive contains raw request measurements, per-example predictions and scores, cache and routing statistics, kernel checks, configuration files, execution logs, and the executed runtime.
  • SHA-256 checksums and archived artifacts link the aggregate tables and figures to the measurements.

8 Ethical Considerations

ExactMoE lowers accelerator-memory requirements for inference, potentially broadening access to sparse language models while also lowering barriers to misuse. It does not change model safety behavior, training data, or licensing obligations, so deployers remain responsible for legal, privacy, content-risk, and application-specific evaluation concerns.

  • ExactMoE reduces accelerator-memory requirements for inference and can improve access to sparse language models.
  • Lower deployment barriers may also broaden misuse of capable models.
  • ExactMoE leaves model safety behavior, training data, and licensing obligations unchanged, while deployers retain responsibility for compliance, privacy, content risks, and application-specific evaluation.

9 Conclusion

ExactMoE combines expert-only W4A16 quantization, kernel-native host storage, configurable GPU residency, and fused grouped execution while retaining all experts. On OLMoE, it substantially reduces memory use while preserving throughput and 99.23% relative normalized-accuracy retention.

  • Memory and throughput: 87.04% lower peak reserved GPU memory was achieved with a 16-expert cache while retaining 81.85% of BF16 decode throughput.The result uses ExactMoE’s configurable GPU residency on OLMoE.
  • Memory and throughput: 47.4% faster full residency was achieved with 71.3% lower reserved memory.This compares full-residency ExactMoE with the BF16 configuration on OLMoE.
  • Accuracy: 99.23% relative normalized-accuracy retention was measured, alongside a small statistically detectable absolute reduction.The accuracy result indicates near-baseline performance despite quantization.
Loading 2608.15383v1…